add mat2json tools (open-mmlab#126)

docs/data_preparation.md

@@ -72,6 +72,15 @@ mmpose
 
 ```
 
+During training and inference, the prediction result will be saved as '.mat' format by default. We also provide a tool to convert this '.mat' to more readable '.json' format.
+```shell
+python tools/mat2json ${PRED_MAT_FILE} ${GT_JSON_FILE} ${OUTPUT_PRED_JSON_FILE}
+```
+For example,
+```shell
+python tools/mat2json work_dirs/res50_mpii_256x256/pred.mat data/mpii/annotations/mpii_val.json pred.json
+```
+
 ## MPII-TRB
 
 For MPII-TRB data, please download from [MPII Human Pose Dataset](http://human-pose.mpi-inf.mpg.de/).

tools/mat2json.py

@@ -0,0 +1,59 @@
+import argparse
+import json
+import time
+
+from scipy.io import loadmat
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description='Converting the predicted .mat file to .json file.')
+    parser.add_argument('pred_mat_file', help='input prediction mat file.')
+    parser.add_argument(
+        'gt_json_file',
+        help='input ground-truth json file to get the image name. '
+        'Default: "data/mpii/mpii_val.json" ')
+    parser.add_argument('output_json_file', help='output converted json file.')
+    args = parser.parse_args()
+    return args
+
+
+def save_json(list_file, path):
+    with open(path, 'w') as f:
+        json.dump(list_file, f, indent=4)
+    return 0
+
+
+def convert_mat(pred_mat_file, gt_json_file, output_json_file):
+    res = loadmat(pred_mat_file)
+    preds = res['preds']
+    N = preds.shape[0]
+
+    with open(gt_json_file) as anno_file:
+        anno = json.load(anno_file)
+
+    assert len(anno) == N
+
+    instance = {}
+
+    for pred, ann in zip(preds, anno):
+        ann.pop('joints_vis')
+        ann['joints'] = pred.tolist()
+
+    instance['annotations'] = anno
+    instance['info'] = {}
+    instance['info']['description'] = 'Converted MPII prediction.'
+    instance['info']['year'] = time.strftime('%Y', time.localtime())
+    instance['info']['date_created'] = time.strftime('%Y/%m/%d',
+                                                     time.localtime())
+
+    save_json(instance, output_json_file)
+
+
+def main():
+    args = parse_args()
+    convert_mat(args.pred_mat_file, args.gt_json_file, args.output_json_file)
+
+
+if __name__ == '__main__':
+    main()