Spaces:
Sleeping
Sleeping
# Copyright (c) OpenMMLab. All rights reserved. | |
import argparse | |
import json | |
import mmengine | |
from mmocr.utils import list_to_file | |
def parse_coco_json(in_path): | |
json_obj = mmengine.load(in_path) | |
image_infos = json_obj['images'] | |
annotations = json_obj['annotations'] | |
imgid2imgname = {} | |
img_ids = [] | |
for image_info in image_infos: | |
imgid2imgname[image_info['id']] = image_info | |
img_ids.append(image_info['id']) | |
imgid2anno = {} | |
for img_id in img_ids: | |
imgid2anno[img_id] = [] | |
for anno in annotations: | |
img_id = anno['image_id'] | |
new_anno = {} | |
new_anno['iscrowd'] = anno['iscrowd'] | |
new_anno['category_id'] = anno['category_id'] | |
new_anno['bbox'] = anno['bbox'] | |
new_anno['segmentation'] = anno['segmentation'] | |
if img_id in imgid2anno.keys(): | |
imgid2anno[img_id].append(new_anno) | |
return imgid2imgname, imgid2anno | |
def gen_line_dict_file(out_path, imgid2imgname, imgid2anno): | |
lines = [] | |
for key, value in imgid2imgname.items(): | |
if key in imgid2anno: | |
anno = imgid2anno[key] | |
line_dict = {} | |
line_dict['file_name'] = value['file_name'] | |
line_dict['height'] = value['height'] | |
line_dict['width'] = value['width'] | |
line_dict['annotations'] = anno | |
lines.append(json.dumps(line_dict)) | |
list_to_file(out_path, lines) | |
def parse_args(): | |
parser = argparse.ArgumentParser() | |
parser.add_argument('--in-path', help='input json path with coco format') | |
parser.add_argument( | |
'--out-path', help='output txt path with line-json format') | |
args = parser.parse_args() | |
return args | |
def main(): | |
args = parse_args() | |
imgid2imgname, imgid2anno = parse_coco_json(args.in_path) | |
gen_line_dict_file(args.out_path, imgid2imgname, imgid2anno) | |
print('finish') | |
if __name__ == '__main__': | |
main() | |