94 lines
2.9 KiB
Python
94 lines
2.9 KiB
Python
import argparse
|
|
from coco.storage import COCOStorage
|
|
import logging
|
|
import pycocotools.coco
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger("tools")
|
|
|
|
|
|
def create(args):
|
|
con = args.storage.con
|
|
cur = con.cursor()
|
|
cur.executemany(
|
|
'INSERT OR IGNORE INTO categories(id, supercategory, name) VALUES (:id, :supercategory, :name)', args.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)
|
|
''', args.coco.imgs.values())
|
|
con.commit()
|
|
|
|
logger.info("Annotations...")
|
|
|
|
def annotation_generator():
|
|
for c in args.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 args.coco.anns.values():
|
|
for i, seg in enumerate(ann['segmentation']):
|
|
yield {
|
|
# create a uniqe segment id, supports max 10 segments per annotation
|
|
'id': ann['id']*10 + i,
|
|
'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...")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
subparsers = parser.add_subparsers(
|
|
title='subcommands', help="Use command -h for specific help")
|
|
|
|
parser_create = subparsers.add_parser('create')
|
|
parser_create.add_argument(
|
|
'--annotations',
|
|
metavar="JSON_FILENAME",
|
|
type=pycocotools.coco.COCO,
|
|
dest='coco',
|
|
default='dataset/annotations/instances_val2017.json'
|
|
)
|
|
parser_create.add_argument(
|
|
'--db',
|
|
type=COCOStorage,
|
|
metavar='DATABASE',
|
|
dest='storage',
|
|
help='SQLite db filename, will be created if not existing',
|
|
default='dataset/instances_val2017.db'
|
|
)
|
|
parser_create.set_defaults(target=create)
|
|
|
|
# parser_build = subparsers.add_parser('build')
|
|
|
|
args = parser.parse_args()
|
|
if 'target' in args:
|
|
args.target(args)
|
|
else:
|
|
parser.print_help()
|
|
exit(1)
|