File size: 1,674 Bytes
1da285f | 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 | 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)
|