37 lines
999 B
Python
37 lines
999 B
Python
|
import pycocotools.coco
|
||
|
|
import argparse
|
||
|
|
import logging
|
||
|
|
import tqdm
|
||
|
|
import urllib.request
|
||
|
|
import os
|
||
|
|
|
||
|
|
logger = logging.getLogger("coco")
|
||
|
|
|
||
|
|
argParser = argparse.ArgumentParser(description='Find all images with single segments and generate svg')
|
||
|
|
argParser.add_argument(
|
||
|
|
'--annotations',
|
||
|
|
type=str,
|
||
|
|
default='../../datasets/COCO/annotations/instances_train2017.json'
|
||
|
|
)
|
||
|
|
argParser.add_argument(
|
||
|
|
'--output',
|
||
|
|
type=str,
|
||
|
|
help='Output directory'
|
||
|
|
)
|
||
|
|
args = argParser.parse_args()
|
||
|
|
|
||
|
|
|
||
|
|
logger.info(f"Load {args.annotations}")
|
||
|
|
coco = pycocotools.coco.COCO(args.annotations)
|
||
|
|
|
||
|
|
|
||
|
|
for img_id, annotations in tqdm.tqdm(coco.imgToAnns.items()):
|
||
|
|
if len(annotations) != 1:
|
||
|
|
continue
|
||
|
|
|
||
|
|
annotation = annotations[0] # we have only one
|
||
|
|
category = coco.cats[annotation['category_id']]['name']
|
||
|
|
fn = os.path.join(args.output, f"{category}_{img_id}.jpg")
|
||
|
|
if not os.path.exists(fn):
|
||
|
|
urllib.request.urlretrieve(coco.imgs[img_id]['coco_url'], fn)
|