mirror of https://github.com/open-mmlab/mmpose
61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
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()
|