Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,071 Bytes
93b49a4 |
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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import json
import os
import shutil
import sys
import numpy as np
from xtcocotools.coco import COCO
def search_match(bbox, num_keypoints, segmentation):
found = []
checked = 0
for json_file, coco in COCO_DICT.items():
cat_ids = coco.getCatIds()
for cat_id in cat_ids:
img_ids = coco.getImgIds(catIds=cat_id)
for img_id in img_ids:
annotations = coco.loadAnns(coco.getAnnIds(imgIds=img_id, catIds=cat_id))
for ann in annotations:
checked += 1
if (ann['num_keypoints'] == num_keypoints and ann['bbox'] == bbox and ann[
'segmentation'] == segmentation):
src_file = coco.loadImgs(img_id)[0]["file_name"]
split = "test" if "test" in json_file else "train"
found.append((src_file, ann, split))
# return src_file, ann, split
if len(found) == 0:
raise Exception("No match found out of {} images".format(checked))
elif len(found) > 1:
raise Exception("More than one match! ".format(found))
return found[0]
if __name__ == "__main__":
carfusion_dir_path = sys.argv[1]
mp100_dataset_path = sys.argv[2]
os.makedirs('output', exist_ok=True)
for cat in ['car', 'bus', 'suv']:
os.makedirs(os.path.join('output', cat), exist_ok=True)
COCO_DICT = {}
ann_files = os.path.join(carfusion_dir_path, 'annotations')
for json_file in os.listdir(ann_files):
COCO_DICT[json_file] = COCO(os.path.join(carfusion_dir_path, 'annotations', json_file))
count = 0
print_log = []
for json_file in os.listdir(mp100_dataset_path):
print("Processing {}".format(json_file))
cats = {}
coco = COCO(os.path.join(mp100_dataset_path, json_file))
cat_ids = coco.getCatIds()
for cat_id in cat_ids:
category_info = coco.loadCats(cat_id)
cat_name = category_info[0]['name']
if cat_name in ['car', 'bus', 'suv']:
cats[cat_name] = cat_id
for cat_name, cat_id in cats.items():
img_ids = coco.getImgIds(catIds=cat_id)
count += len(img_ids)
print_log.append(f'{json_file} : {cat_name}: {len(img_ids)}')
for img_id in img_ids:
img = coco.loadImgs(img_id)[0]
dst_file_name = img['file_name']
annotation = coco.loadAnns(coco.getAnnIds(imgIds=img_id, catIds=cat_id, iscrowd=None))
bbox = annotation[0]['bbox']
keypoints = annotation[0]['keypoints']
segmentation = annotation[0]['segmentation']
num_keypoints = annotation[0]['num_keypoints']
# Search for a match:
src_img, src_ann, split = search_match(bbox, num_keypoints, segmentation)
shutil.copyfile(
os.path.join(carfusion_dir_path, split, src_img),
os.path.join('output', dst_file_name)) |