mirror of https://github.com/RE-OWOD/RE-OWOD
64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
|
from pycocotools.coco import COCO
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
T2_CLASS_NAMES = [
|
||
|
|
"truck", "traffic light", "fire hydrant", "stop sign", "parking meter",
|
||
|
|
"bench", "elephant", "bear", "zebra", "giraffe",
|
||
|
|
"backpack", "umbrella", "handbag", "tie", "suitcase",
|
||
|
|
"microwave", "oven", "toaster", "sink", "refrigerator"
|
||
|
|
]
|
||
|
|
|
||
|
|
# Train
|
||
|
|
coco_annotation_file = '/home/joseph/workspace/datasets/mscoco/annotations/instances_train2017.json'
|
||
|
|
dest_file = '/home/joseph/workspace/OWOD/datasets/coco17_voc_style/ImageSets/t2_train.txt'
|
||
|
|
|
||
|
|
coco_instance = COCO(coco_annotation_file)
|
||
|
|
|
||
|
|
image_ids = []
|
||
|
|
cls = []
|
||
|
|
for index, image_id in enumerate(coco_instance.imgToAnns):
|
||
|
|
image_details = coco_instance.imgs[image_id]
|
||
|
|
classes = [coco_instance.cats[annotation['category_id']]['name'] for annotation in coco_instance.imgToAnns[image_id]]
|
||
|
|
if not set(classes).isdisjoint(T2_CLASS_NAMES):
|
||
|
|
image_ids.append(image_details['file_name'].split('.')[0])
|
||
|
|
cls.extend(classes)
|
||
|
|
|
||
|
|
(unique, counts) = np.unique(cls, return_counts=True)
|
||
|
|
print({x:y for x,y in zip(unique, counts)})
|
||
|
|
|
||
|
|
with open(dest_file, 'w') as file:
|
||
|
|
for image_id in image_ids:
|
||
|
|
file.write(str(image_id)+'\n')
|
||
|
|
|
||
|
|
print('Created train file')
|
||
|
|
|
||
|
|
# Test
|
||
|
|
coco_annotation_file = '/home/joseph/workspace/datasets/mscoco/annotations/instances_val2017.json'
|
||
|
|
dest_file = '/home/joseph/workspace/OWOD/datasets/coco17_voc_style/ImageSets/t2_test.txt'
|
||
|
|
|
||
|
|
coco_instance = COCO(coco_annotation_file)
|
||
|
|
|
||
|
|
image_ids = []
|
||
|
|
cls = []
|
||
|
|
for index, image_id in enumerate(coco_instance.imgToAnns):
|
||
|
|
image_details = coco_instance.imgs[image_id]
|
||
|
|
classes = [coco_instance.cats[annotation['category_id']]['name'] for annotation in coco_instance.imgToAnns[image_id]]
|
||
|
|
if not set(classes).isdisjoint(T2_CLASS_NAMES):
|
||
|
|
image_ids.append(image_details['file_name'].split('.')[0])
|
||
|
|
cls.extend(classes)
|
||
|
|
|
||
|
|
(unique, counts) = np.unique(cls, return_counts=True)
|
||
|
|
print({x:y for x,y in zip(unique, counts)})
|
||
|
|
|
||
|
|
with open(dest_file, 'w') as file:
|
||
|
|
for image_id in image_ids:
|
||
|
|
file.write(str(image_id)+'\n')
|
||
|
|
print('Created test file')
|
||
|
|
|
||
|
|
dest_file = '/home/joseph/workspace/OWOD/datasets/coco17_voc_style/ImageSets/t2_test_unk.txt'
|
||
|
|
with open(dest_file, 'w') as file:
|
||
|
|
for image_id in image_ids:
|
||
|
|
file.write(str(image_id)+'\n')
|
||
|
|
|
||
|
|
print('Created test_unk file')
|