106 lines
3.2 KiB
Python
106 lines
3.2 KiB
Python
|
import pycocotools.coco
|
||
|
|
import argparse
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
import pprint
|
||
|
|
import sqlite3
|
||
|
|
|
||
|
|
logging.basicConfig(level=logging.INFO)
|
||
|
|
logger = logging.getLogger("coco")
|
||
|
|
|
||
|
|
argParser = argparse.ArgumentParser(description='Create shape SVG\'s')
|
||
|
|
argParser.add_argument(
|
||
|
|
'--annotations',
|
||
|
|
type=str,
|
||
|
|
default='../../datasets/COCO/annotations/instances_val2017.json'
|
||
|
|
)
|
||
|
|
argParser.add_argument(
|
||
|
|
'--categories',
|
||
|
|
action='store_true',
|
||
|
|
help='Show categories'
|
||
|
|
)
|
||
|
|
argParser.add_argument(
|
||
|
|
'--propagate',
|
||
|
|
type=str,
|
||
|
|
metavar='DATABASE',
|
||
|
|
help='Store data in sqlite db'
|
||
|
|
)
|
||
|
|
args = argParser.parse_args()
|
||
|
|
|
||
|
|
|
||
|
|
logger.info(f"Load {args.annotations}")
|
||
|
|
coco = pycocotools.coco.COCO(args.annotations)
|
||
|
|
|
||
|
|
|
||
|
|
if args.categories:
|
||
|
|
cats = {}
|
||
|
|
for id, cat in coco.cats.items():
|
||
|
|
if cat['supercategory'] not in cats:
|
||
|
|
cats[cat['supercategory']] = []
|
||
|
|
cats[cat['supercategory']].append(cat)
|
||
|
|
# pp = pprint.PrettyPrinter(indent=4)
|
||
|
|
pprint.pprint(cats, sort_dicts=False)
|
||
|
|
|
||
|
|
if args.propagate:
|
||
|
|
if not os.path.exists(args.propagate):
|
||
|
|
con = sqlite3.connect(args.propagate)
|
||
|
|
cur = con.cursor()
|
||
|
|
with open('coco.sql', 'r') as fp:
|
||
|
|
cur.executescript(fp.read())
|
||
|
|
con.close()
|
||
|
|
|
||
|
|
con = sqlite3.connect(args.propagate)
|
||
|
|
logger.info("Create categories")
|
||
|
|
cur = con.cursor()
|
||
|
|
cur.executemany('INSERT OR IGNORE INTO categories(id, supercategory, name) VALUES (:id, :supercategory, :name)', coco.cats.values())
|
||
|
|
con.commit()
|
||
|
|
|
||
|
|
logger.info("Images...")
|
||
|
|
cur.executemany('''
|
||
|
|
INSERT OR IGNORE INTO images(id, flickr_url, coco_url, width, height, date_captured)
|
||
|
|
VALUES (:id, :flickr_url, :coco_url, :width, :height, :date_captured)
|
||
|
|
''', coco.imgs.values())
|
||
|
|
con.commit()
|
||
|
|
|
||
|
|
logger.info("Annotations...")
|
||
|
|
|
||
|
|
|
||
|
|
def annotation_generator():
|
||
|
|
for c in coco.anns.values():
|
||
|
|
ann = c.copy()
|
||
|
|
ann['bbox_top'] = ann['bbox'][1]
|
||
|
|
ann['bbox_left'] = ann['bbox'][0]
|
||
|
|
ann['bbox_width'] = ann['bbox'][2]
|
||
|
|
ann['bbox_height'] = ann['bbox'][3]
|
||
|
|
yield ann
|
||
|
|
|
||
|
|
cur.executemany('''
|
||
|
|
INSERT OR IGNORE INTO annotations(id, image_id, category_id, iscrowd, area, bbox_top, bbox_left, bbox_width, bbox_height)
|
||
|
|
VALUES (:id, :image_id, :category_id, :iscrowd, :area, :bbox_top, :bbox_left, :bbox_width, :bbox_height)
|
||
|
|
''', annotation_generator())
|
||
|
|
con.commit()
|
||
|
|
|
||
|
|
|
||
|
|
logger.info("Segments...")
|
||
|
|
|
||
|
|
def segment_generator():
|
||
|
|
for ann in coco.anns.values():
|
||
|
|
for i, seg in enumerate(ann['segmentation']):
|
||
|
|
yield {
|
||
|
|
'id': ann['id']*10 + i, # create a uniqe segment id, supports max 10 segments per annotation
|
||
|
|
'annotation_id': ann['id'],
|
||
|
|
'points': str(seg)[1:-1],
|
||
|
|
}
|
||
|
|
|
||
|
|
cur.executemany('''
|
||
|
|
INSERT OR IGNORE INTO segments(id, annotation_id, points)
|
||
|
|
VALUES (:id, :annotation_id, :points)
|
||
|
|
''', segment_generator())
|
||
|
|
con.commit()
|
||
|
|
|
||
|
|
|
||
|
|
logger.info("Done...")
|
||
|
|
|
||
|
|
# for id, cat in coco.cats.items():
|
||
|
|
# cur = con.cursor()
|
||
|
|
# cur.execute
|