|
|
import json
|
|
|
import pickle
|
|
|
import os
|
|
|
from tqdm import tqdm
|
|
|
|
|
|
def build_referring_dataset(instance_path, refs_path, split, save_path):
|
|
|
assert os.path.exists(instance_path), f'Path not found: {instance_path}'
|
|
|
assert os.path.exists(refs_path), f'Path not found: {refs_path}'
|
|
|
|
|
|
with open(instance_path) as f:
|
|
|
instance = json.load(f)
|
|
|
with open(refs_path, 'rb') as f:
|
|
|
refs = pickle.load(f)
|
|
|
|
|
|
images = instance['images']
|
|
|
annotations = instance['annotations']
|
|
|
|
|
|
img_id2info = {}
|
|
|
for image in images:
|
|
|
img_id2info[image['id']] = image
|
|
|
anno_id2info = {}
|
|
|
for annotation in annotations:
|
|
|
anno_id2info[annotation['id']] = annotation
|
|
|
|
|
|
outputs = []
|
|
|
new_img_id = 0
|
|
|
for sample in tqdm(refs):
|
|
|
if sample['split'] != split:
|
|
|
continue
|
|
|
if -1 in sample['ann_id']:
|
|
|
image = sample['file_name']
|
|
|
image_id = sample['image_id']
|
|
|
image_info = img_id2info[image_id]
|
|
|
instruction = sample['sentences']
|
|
|
ann_ids = sample['ann_id']
|
|
|
anns = []
|
|
|
result = {
|
|
|
'image': image,
|
|
|
'image_info': image_info,
|
|
|
'instruction': instruction,
|
|
|
'new_img_id': new_img_id,
|
|
|
'anns': anns
|
|
|
}
|
|
|
outputs.append(result)
|
|
|
new_img_id += 1
|
|
|
continue
|
|
|
image = sample['file_name']
|
|
|
image_id = sample['image_id']
|
|
|
image_info = img_id2info[image_id]
|
|
|
instruction = sample['sentences']
|
|
|
ann_ids = sample['ann_id']
|
|
|
anns = [anno_id2info[id] for id in ann_ids]
|
|
|
result = {
|
|
|
'image': image,
|
|
|
'image_info': image_info,
|
|
|
'instruction': instruction,
|
|
|
'new_img_id': new_img_id,
|
|
|
'anns': anns
|
|
|
}
|
|
|
outputs.append(result)
|
|
|
new_img_id += 1
|
|
|
|
|
|
with open(save_path, 'w') as f:
|
|
|
json.dump(outputs, f)
|
|
|
print(f'Saving at {save_path}. Total sample: {len(outputs)}.')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
|
|
root_path = 'datasets/refer_seg'
|
|
|
datasets = 'grefcoco'
|
|
|
splits = ['train', 'val', 'testA', 'testB']
|
|
|
|
|
|
for split in splits:
|
|
|
instance_path = os.path.join(root_path, datasets, 'instances.json')
|
|
|
refs_path = os.path.join(root_path, datasets, 'grefs(unc).json')
|
|
|
save_path = os.path.join(root_path, datasets, f'{split}_psalm.json')
|
|
|
print(f'Processing gRefCOCO: {split}...')
|
|
|
|
|
|
build_referring_dataset(instance_path, refs_path, split, save_path) |
