auto-accept/tools.py

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)