auto-accept/coco/storage.py

import logging
import os
import sqlite3
import pycocotools.coco
import ast
import copy
import svgwrite
import uuid

logger = logging.getLogger('coco.storage')


class Annotation:
    def __init__(self, result, storage):
        self.storage = storage
        self.id = result['id']
        self.image_id = result['image_id']
        self.category_id = result['category_id']
        self.iscrowd = bool(result['iscrowd'])
        self.area = result['area']
        self.bbox = [result['bbox_left'], result['bbox_top'],
                     result['bbox_width'], result['bbox_height']]
        self.segments = self.fetchSegments()
        self.is_normalised = False
        if 'zerkine_moment' in result and type(result['zerkine_moment']) is list:
            # when normalising, this is already there
            self.zerkine_moment = result['zerkine_moment']
        else:
            self.zerkine_moment = self.parseZerkineFromDB(
                result['zerkine_moment']) if 'zerkine_moment' in result else None

    @classmethod
    def parseZerkineFromDB(cls, r):
        z = r.split(' ')
        return [float(i) for i in z]

    def fetchSegments(self):
        try:
            cur = self.storage.con.cursor()
            cur.execute(
                "SELECT * FROM segments WHERE annotation_id = :id AND points != 'ount' AND points != 'iz'", {'id': self.id})
            segments = []
            for row in cur:
                segments.append(Segment(row))
        except Exception as e:
            logger.critical(f"Invalid segment for annotation {self.id}")
            logger.exception(e)
            raise(e)
        return segments

    def getNormalised(self, width, height) -> 'Annotation':
        '''
        center segments in boundig box with given width and height, and on point 0,0
        '''
        scale = min(width/self.bbox[2], height/self.bbox[3])
        logger.debug(f"Normalise from bbox: {self.bbox}")
        new_width = self.bbox[2] * scale
        new_height = self.bbox[3] * scale

        dx = (width - new_width) / 2
        dy = (height - new_height) / 2

        data = self.forJson()
        data['bbox_left'] = 0
        data['bbox_top'] = 0
        data['bbox_width'] = new_width
        data['bbox_height'] = new_height
        newAnn = Annotation(data, self.storage)
        newAnn.is_normalised = True
        newAnn.bbox_original = self.bbox
        newAnn.scale = scale

        for i, segment in enumerate(newAnn.segments):
            newAnn.segments[i].points = [[
                (p[0]-self.bbox[0]) * scale,
                (p[1]-self.bbox[1]) * scale
            ] for p in segment.points]

        return newAnn

    def forJson(self):
        data = self.__dict__.copy()
        del data['storage']
        data['image'] = self.storage.getImage(data['image_id'])
        data['category'] = self.storage.getCategory(self.category_id)
        return data

    def writeToDrawing(self, dwg, **pathSpecs):
        for segment in self.segments:
            if len(pathSpecs) == 0:
                pathSpecs['fill'] = 'white'
            dwg.add(svgwrite.path.Path(segment.getD(),
                                       class_=f"cat_{self.category_id}", **pathSpecs))

    def getTranslationToCenter(self):
        dimensions = (self.bbox[2], self.bbox[3])
        targetSize = max(dimensions)
        dx = (dimensions[0] - targetSize)/2
        dy = (dimensions[1] - targetSize)/2
        return (dx, dy)

    def asSvg(self, filename, square=False, bg=None, clip_image=False, image_dir=None) -> svgwrite.Drawing:
        dimensions = (self.bbox[2], self.bbox[3])
        viewbox = copy.copy(self.bbox)
        if square:
            targetSize = max(dimensions)
            dx = (dimensions[0] - targetSize)/2
            dy = (dimensions[1] - targetSize)/2
            viewbox[2] = targetSize
            viewbox[3] = targetSize
            dimensions = (targetSize, targetSize)
            viewbox[0] += dx
            viewbox[1] += dy
        dwg = svgwrite.Drawing(
            filename,
            size=dimensions,
            viewBox=" ".join([str(s) for s in viewbox])
        )

        if bg:
            dwg.add(dwg.rect(
                (viewbox[0], viewbox[1]),
                (viewbox[2], viewbox[3]),
                fill=bg))
        if clip_image is not False:
            clip_id = uuid.uuid4().hex

            img = self.storage.getImage(self.image_id)
            rmpart = len("http://images.cocodataset.org/")
            href = img['coco_url'][rmpart:]
            if image_dir:
                href = os.path.join(image_dir, href)
            dwg.add(dwg.image(href, insert=(0,0), clip_path=f"url(#{clip_id})"))

            
            clip_path = dwg.defs.add(dwg.clipPath(id=clip_id))
            self.writeToDrawing(clip_path)
        else:
            self.writeToDrawing(dwg)
        return dwg


class Segment():
    def __init__(self, result):
        try:
            self.points = self.asCoordinates(
                ast.literal_eval('['+result['points']+']'))
        except Exception as e:
            logger.critical(
                f"Exception loading segment for {result} {result['points']}")
            raise

    @classmethod
    def asCoordinates(cls, pointList):
        points = []

        r = len(pointList) / 2
        for i in range(int(r)):
            points.append([
                pointList[(i)*2],
                pointList[(i)*2+1]
            ])
        return points

    def getD(self):
        start = self.points[0]
        d = f'M{start[0]:.4f} {start[1]:.4f} L'
        for i in range(1, len(self.points)):
            p = self.points[i]
            d += f' {p[0]:.4f} {p[1]:.4f}'
        d += " Z"  # segments are always closed
        return d

    def forJson(self):
        return self.points


class COCOStorage:
    def __init__(self, filename):
        self.logger = logging.getLogger('coco.storage')
        self.filename = filename
        if not os.path.exists(self.filename):
            con = sqlite3.connect(self.filename)
            cur = con.cursor()
            d = os.path.dirname(os.path.realpath(__file__))
            with open(os.path.join(d, 'coco.sql'), 'r') as fp:
                cur.executescript(fp.read())
            con.close()

        self.con = sqlite3.connect(self.filename)
        self.con.row_factory = sqlite3.Row

    def propagateFromAnnotations(self, filename):
        self.logger.info(f"Load {filename}")
        coco = pycocotools.coco.COCO(filename)

        self.logger.info("Create categories")
        cur = self.con.cursor()
        cur.executemany(
            'INSERT OR IGNORE INTO categories(id, supercategory, name) VALUES (:id, :supercategory, :name)', coco.cats.values())
        self.con.commit()

        self.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())
        self.con.commit()

        self.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())
        self.con.commit()

        self.logger.info("Segments...")

        def segment_generator():
            for ann in 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())
        self.con.commit()

        self.logger.info("Done...")

    def getCategories(self):
        if not hasattr(self, 'categories'):
            cur = self.con.cursor()
            cur.execute("SELECT * FROM categories ORDER BY id")
            self.categories = [dict(cat) for cat in cur]
        return self.categories

    def getCategory(self, cid):
        cats = self.getCategories()
        cat = [c for c in cats if c['id'] == cid]
        if not len(cat):
            return None
        return cat[0]

    def getImage(self, image_id: int):
        cur = self.con.cursor()
        cur.execute(f"SELECT * FROM images WHERE id = ? LIMIT 1", (image_id,))
        img = cur.fetchone()
        return dict(img)

    def getAnnotationWithoutZerkine(self):
        cur = self.con.cursor()
        # annotation 918 and 2206849 have 0 height. Crashing the script... exclude them
        cur.execute(
            f"SELECT * FROM annotations WHERE zerkine_moment IS NULL AND area > 0 LIMIT 1")
        ann = cur.fetchone()
        if ann:
            return Annotation(ann, self)
        else:
            return None

    def countAnnotationsWithoutZerkine(self):
        cur = self.con.cursor()

        cur.execute(
            f"SELECT count(id) FROM annotations WHERE zerkine_moment IS NULL AND area > 0")
        return int(cur.fetchone()[0])

    def storeZerkineForAnnotation(self, annotation, moments, delayCommit=False):
        m = ' '.join([str(m) for m in moments])
        cur = self.con.cursor()

        cur.execute(
            "UPDATE annotations SET zerkine_moment = :z WHERE id = :id",
            {'z': m, 'id': annotation.id}
        )
        if not delayCommit:
            self.con.commit()
        return True

    def getZerkines(self):
        cur = self.con.cursor()
        cur.execute(
            f"SELECT id, zerkine_moment FROM annotations WHERE zerkine_moment IS NOT NULL")
        return cur.fetchall()

    def getAllAnnotationPoints(self):
        cur = self.con.cursor()
        cur.execute(
            f"SELECT annotations.id, points FROM annotations INNER JOIN segments ON segments.annotation_id = annotations.id WHERE area > 0")
        return cur.fetchall()

    def getAnnotationById(self, annotation_id=None, withZerkine=False) -> Annotation:
        if annotation_id == -1:
            annotation_id = None
        return self.getRandomAnnotation(annotation_id=annotation_id, withZerkine=withZerkine)

    def getRandomAnnotation(self, annotation_id=None, category_id=None, withZerkine=False, min_area=None) -> Annotation:
        result = self.getRandomAnnotations(
            annotation_id, category_id, withZerkine, min_area=min_area, limit=1)
        return result[0] if len(result) else None

    def getRandomAnnotations(self, annotation_id=None, category_id=None, withZerkine=False, min_area=None, limit=None):
        cur = self.con.cursor()
        where = ""
        params = []
        if annotation_id:
            where = "id = ?"
            params.append(annotation_id)
        elif category_id:
            where = "category_id = ?"
            params.append(category_id)
        else:
            where = "1=1"

        if withZerkine:
            where += " AND zerkine_moment IS NOT NULL"

        if min_area:
            where += " AND area > ?"
            params.append(min_area)
            print(where)
        

        sqlLimit = ""
        if limit:
            sqlLimit = f"LIMIT {int(limit)}"

        cur.execute(
            f"SELECT * FROM annotations WHERE {where} ORDER BY RANDOM() {sqlLimit}", tuple(params))
        results = []
        for ann in cur:
            results.append(Annotation(ann, self))
        return results
#             ann = cur.fetchall()
#
#             if ann:
#                 return Annotation(ann, self)
#             else:
#                 return None