mmpose/tools/dataset_converters/mat2json.py

# 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()