File size: 1,213 Bytes
814a594 |
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 |
import os
import json
from refer import REFER
coco_root = '/pth/to/coco'
ref_root = '/pth/to/refcocoseg'
coco_train_annot = json.load(open(os.path.join(coco_root, 'annotations/instances_train2017.json')))
coco_train_id = []
image_annot = {}
for i in range(len(coco_train_annot['images'])):
coco_train_id.append(coco_train_annot['images'][i]['id'])
image_annot[coco_train_annot['images'][i]['id']] = coco_train_annot['images'][i]
refg = REFER(data_root=ref_root,
dataset='refcocog', splitBy='umd')
refg_val_ids = refg.getRefIds(split='val')
full_anno = []
for ref_id in refg_val_ids:
ref = refg.loadRefs(ref_id)[0]
anno = refg.refToAnn[ref_id]
anno.update(ref)
full_anno.append(anno)
imageid_list = []
final_anno = {}
for anno in full_anno:
imageid_list += [anno['image_id']]
final_anno[anno['ann_id']] = anno
annotations = [value for key, value in final_anno.items()]
iamges = []
for image_id in list(set(imageid_list)):
iamges += [image_annot[image_id]]
outputs = {'images': iamges, 'annotations': annotations}
print(len(iamges))
print(len(annotations))
json.dump(outputs, open(os.path.join(coco_root, 'annotations/refcocog_umd_train.json'), 'w'))
|