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)