auto-accept/tools.py

95 lines
2.9 KiB
Python
Raw Normal View History

2020-12-17 14:57:03 +01:00
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()
2020-12-17 15:22:24 +01:00
cur.executemany(
'INSERT OR IGNORE INTO categories(id, supercategory, name) VALUES (:id, :supercategory, :name)', args.coco.cats.values())
2020-12-17 14:57:03 +01:00
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 {
2020-12-17 15:22:24 +01:00
# create a uniqe segment id, supports max 10 segments per annotation
'id': ann['id']*10 + i,
2020-12-17 14:57:03 +01:00
'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__":
2020-12-17 15:22:24 +01:00
2020-12-17 14:57:03 +01:00
parser = argparse.ArgumentParser()
2020-12-17 15:22:24 +01:00
subparsers = parser.add_subparsers(
title='subcommands', help="Use command -h for specific help")
2020-12-17 14:57:03 +01:00
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'
2020-12-17 15:22:24 +01:00
)
2020-12-17 14:57:03 +01:00
parser_create.add_argument(
2020-12-17 15:22:24 +01:00
'--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)
2020-12-17 14:57:03 +01:00
# parser_build = subparsers.add_parser('build')
args = parser.parse_args()
if 'target' in args:
args.target(args)
else:
parser.print_help()
2020-12-17 15:22:24 +01:00
exit(1)