| import os |
| import argparse |
| import json |
| from collections import defaultdict |
|
|
|
|
| def uied2coco(uied_dir, img_id_map, cat_id_map, output_file = None): |
| pred = [] |
| for file in os.listdir(uied_dir): |
| if not file.endswith('.json'): |
| continue |
| img_id = img_id_map[file.split('.')[-2]] |
| if img_id is None: |
| continue |
| with open(os.path.join(uied_dir, file), 'r') as f: |
| data = json.load(f)['compos'] |
| for i in range(len(data)): |
| cat_id = cat_id_map[data[i]['class']] |
| x = data[i]['position']['column_min'] |
| y = data[i]['position']['row_min'] |
| width = data[i]['width'] |
| height = data[i]['height'] |
| bbox = [x, y, width, height] |
| score = 1 |
| pred.append({'image_id': img_id, 'category_id': cat_id, 'bbox': bbox, 'score': score}) |
| if output_file is not None: |
| with open(output_file, 'w') as f: |
| json.dump(pred, f) |
| return pred |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--uied_dir', type=str, required=True) |
| parser.add_argument('--ann', type=str, required=True) |
| parser.add_argument('--output_file', type=str, required=True) |
| args = parser.parse_args() |
| with open(args.ann, 'r') as f: |
| ann = json.load(f) |
| img_id_map = defaultdict(lambda: None) |
| for img in ann['images']: |
| img_id_map[img['file_name'].split('.')[-2]] = img['id'] |
| cat_id_map = defaultdict(lambda: 1) |
| for cat in ann['categories']: |
| cat_id_map[cat['name']] = cat['id'] |
| uied2coco(args.uied_dir, img_id_map, cat_id_map, args.output_file) |
| |
