File size: 1,649 Bytes
3bbb319 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
# 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()
|