File size: 1,811 Bytes
625a17f |
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 48 49 50 51 52 53 54 55 |
import json
import os
from pycocotools.coco import COCO
from tqdm import tqdm
import concurrent.futures
import re
if __name__ == '__main__':
root_path = 'datasets/lvis'
splits = ['train', 'val']
for split in splits:
print(f'Processing {split}.')
annotation_path = os.path.join(root_path, f'lvis_v1_{split}.json')
save_path = os.path.join(root_path, f'lvis_instance_psalm.json')
with open(annotation_path) as f:
data = json.load(f)
categories = data['categories']
categories_path = os.path.join(root_path, 'lvis_categories.json')
with open(categories_path, 'w') as f:
json.dump(categories, f, indent=2)
annotation_map = {}
for anno in data['annotations']:
image_id = anno['image_id']
if image_id not in annotation_map:
annotation_map[image_id] = []
annotation_map[image_id].append(anno)
lvis_anno = []
pattern = re.compile(r'.*/((?:train|val)\d+/\d+\.jpg)')
new_img_id = 0
for img in tqdm(data['images']):
if img['id'] not in annotation_map:
continue
match = pattern.search(img['coco_url'])
if match:
image = match.group(1)
else:
image = img['coco_url']
lvis_anno.append(
{
'image': image,
'image_info': img,
'new_img_id': new_img_id,
'anns': annotation_map[img['id']]
}
)
new_img_id += 1
with open(save_path, 'w') as f:
json.dump(lvis_anno, f, indent=2)
print(f'saving at {save_path}.')
|