Copy annotations (#55)

* ignore local/generated files

* remove speculative code bug

* working propagate next/previous

* use a separate route for copying annotations

* use single file; using single initializer method

* upsert function

* upsert category test

* validate from image ids

* copy endpoint

* simplified category create

* all images endpoint

* copy annotations

* copy annotation bug fixes

* copy selected categories

* test case fix

Author: matt-deboer

app/api/annotator.py

@@ -1,4 +1,5 @@
-from flask_restplus import Namespace, Api, Resource
+import copy
+from flask_restplus import Namespace, Api, Resource, reqparse
 from flask import request
 
 from ..util import query_util
@@ -95,7 +96,6 @@ class AnnotatorId(Resource):
 
     def get(self, image_id):
         """ Called when loading from the annotator client """
-
         image = ImageModel.objects(id=image_id).first()
 
         if image is None:

app/api/categories.py

@@ -35,12 +35,13 @@ def post(self):
         color = args.get('color')
 
         try:
-            category = CategoryModel.create_category(
+            category = CategoryModel(
                 name=name,
                 supercategory=supercategory,
                 color=color,
                 metadata=metadata
             )
+            category.save()
         except (ValueError, TypeError) as e:
             return {'message': str(e)}, 400
 

app/api/datasets.py

@@ -60,19 +60,7 @@ def post(self):
         name = args['name']
         categories = args.get('categories', [])
 
-        category_ids = []
-
-        for category in categories:
-            if isinstance(category, int):
-                category_ids.append(category)
-            else:
-                category_model = CategoryModel.objects(name=category).first()
-
-                if category_model is None:
-                    new_category = CategoryModel.create_category(name=category)
-                    category_ids.append(new_category.id)
-                else:
-                    category_ids.append(category_model.id)
+        category_ids = CategoryModel.bulk_create(categories)
 
         try:
             dataset = DatasetModel(name=name, categories=category_ids)
@@ -139,21 +127,8 @@ def post(self, dataset_id):
         default_annotation_metadata = args.get('default_annotation_metadata')
 
         if categories is not None:
-            category_ids = []
-
-            for category in categories:
-                if isinstance(category, int):
-                    category_ids.append(category)
-                else:
-                    category_model = CategoryModel.objects(name=category).first()
-
-                    if category_model is None:
-                        new_category = CategoryModel.create_category(name=category)
-                        category_ids.append(new_category.id)
-                    else:
-                        category_ids.append(category_model.id)
+            dataset.categories = CategoryModel.bulk_create(categories)
 
-            dataset.categories = category_ids
 
         if default_annotation_metadata is not None:
             dataset.default_annotation_metadata = default_annotation_metadata

app/api/images.py

@@ -1,6 +1,5 @@
 from flask_restplus import Namespace, Resource, reqparse
 from werkzeug.datastructures import FileStorage
-
 from flask import send_file
 
 from ..util import query_util, coco_util, thumbnail_util
@@ -14,6 +13,12 @@
 api = Namespace('image', description='Image related operations')
 
 
+image_all = reqparse.RequestParser()
+image_all.add_argument('fields', required=False, type=str)
+image_all.add_argument('page', default=1, type=int)
+image_all.add_argument('perPage', default=50, type=int, required=False)
+
+
 image_upload = reqparse.RequestParser()
 image_upload.add_argument('image', location='files',
                           type=FileStorage, required=True,
@@ -26,12 +31,37 @@
 image_download.add_argument('width', type=int, required=False, default=0)
 image_download.add_argument('height', type=int, required=False, default=0)
 
+copy_annotations = reqparse.RequestParser()
+copy_annotations.add_argument('category_ids', location='json', type=list,
+                              required=False, default=None, help='Categories to copy')
+
 
 @api.route('/')
 class Images(Resource):
+    @api.expect(image_all)
     def get(self):
         """ Returns all images """
-        return query_util.fix_ids(ImageModel.objects(deteled=False).all())
+        args = image_all.parse_args()
+        per_page = args['perPage']
+        page = args['page']-1
+        fields = args.get('fields', "")
+
+        images = ImageModel.objects(deleted=False)
+        total = images.count()
+        pages = int(total/per_page) + 1
+
+        images = images.skip(page*per_page).limit(per_page)
+        if fields:
+            images = images.only(*fields.split(','))
+
+        return {
+            "total": total,
+            "pages": pages,
+            "page": page,
+            "fields": fields,
+            "per_page": per_page,
+            "images": query_util.fix_ids(images.all())
+        }
 
     @api.expect(image_upload)
     def post(self):
@@ -114,6 +144,38 @@ def delete(self, image_id):
         return {"success": True}
 
 
+@api.route('/copy/<int:from_id>/<int:to_id>/annotations')
+class ImageCopyAnnotations(Resource):
+
+    @api.expect(copy_annotations)
+    def post(self, from_id, to_id):
+        args = copy_annotations.parse_args()
+        category_ids = args.get('category_ids')
+
+        image_from = ImageModel.objects(id=from_id).first()
+        image_to = ImageModel.objects(id=to_id).first()
+
+        if image_from is None or image_to is None:
+            return {'success': False, 'message': 'Invalid image ids'}, 400
+
+        if image_from == image_to:
+            return {'success': False, 'message': 'Cannot copy self'}, 400
+
+        if image_from.width != image_to.width or image_from.height != image_to.height:
+            return {'success': False, 'message': 'Image sizes do not match'}, 400
+
+        if category_ids is None:
+            category_ids = DatasetModel.objects(id=image_from.dataset_id).first().categories
+
+        query = AnnotationModel.objects(
+            image_id=image_from.id,
+            category_id__in=category_ids,
+            deleted=False
+        )
+
+        return {'annotations_created': image_to.copy_annotations(query)}
+
+
 @api.route('/<int:image_id>/thumbnail')
 class ImageCoco(Resource):
 

app/models.py

@@ -1,12 +1,15 @@
 import os
 import sys
 import json
+import copy
+import numpy as np
+
 from flask_mongoengine import MongoEngine
+from .util.coco_util import decode_seg
 from .util import color_util
 from .config import Config
 from PIL import Image
 
-
 db = MongoEngine()
 
 
@@ -115,6 +118,24 @@ def thumbnail_path(self):
 
         return '/'.join(folders)
 
+    def copy_annotations(self, annotations):
+        """
+        Creates a copy of the annotations for this image
+        :param annotations: QuerySet of annotation models
+        :return: number of annotations
+        """
+        annotations = annotations.filter(width=self.width, height=self.height, area__gt=0)
+
+        for annotation in annotations:
+            clone = annotation.clone()
+
+            clone.dataset_id = self.dataset_id
+            clone.image_id = self.id
+
+            clone.save(copy=True)
+
+        return annotations.count()
+
 
 class AnnotationModel(db.DynamicDocument):
 
@@ -131,7 +152,7 @@ class AnnotationModel(db.DynamicDocument):
     width = db.IntField()
     height = db.IntField()
 
-    color = db.StringField(default=color_util.random_color_hex())
+    color = db.StringField()
 
     metadata = db.DictField(default={})
     paper_object = db.ListField(default=[])
@@ -153,21 +174,34 @@ def __init__(self, image_id=None, **data):
 
         super(AnnotationModel, self).__init__(**data)
 
-    def save(self, *args, **kwargs):
+    def save(self, copy=False, *args, **kwargs):
 
-        if self.dataset_id is not None:
+        if not self.dataset_id and not copy:
             dataset = DatasetModel.objects(id=self.dataset_id).first()
 
             if dataset is not None:
-                metadata = dataset.default_annotation_metadata.copy()
-                metadata.update(self.metadata)
-                self.metadata = metadata
+                self.metadata = dataset.default_annotation_metadata.copy()
+
+        if self.color is None:
+            self.color = color_util.random_color_hex()
 
         return super(AnnotationModel, self).save(*args, **kwargs)
 
     def is_empty(self):
         return len(self.segmentation) == 0 or self.area == 0
 
+    def mask(self):
+        """ Returns binary mask of annotation """
+        mask = np.zeros((self.height, self.width))
+        return decode_seg(mask, self.segmentation)
+
+    def clone(self):
+        """ Creates a clone """
+        create = json.loads(self.to_json())
+        del create['_id']
+
+        return AnnotationModel(**create)
+
 
 class CategoryModel(db.DynamicDocument):
     id = db.SequenceField(primary_key=True)
@@ -180,12 +214,30 @@ class CategoryModel(db.DynamicDocument):
     deleted_date = db.DateTimeField()
 
     @classmethod
-    def create_category(cls, name, color=None, metadata=None, supercategory=None):
-        category = CategoryModel(name=name, supercategory=supercategory)
-        category.metadata = metadata if metadata is not None else {}
-        category.color = color_util.random_color_hex() if color is None else color
-        category.save()
-        return category
+    def bulk_create(cls, categories):
+
+        if not categories:
+            return []
+
+        category_ids = []
+        for category in categories:
+            category_model = CategoryModel.objects(name=category).first()
+
+            if category_model is None:
+                new_category = CategoryModel(name=category)
+                new_category.save()
+                category_ids.append(new_category.id)
+            else:
+                category_ids.append(category_model.id)
+
+        return category_ids
+
+    def save(self, *args, **kwargs):
+
+        if not self.color:
+            self.color = color_util.random_color_hex()
+
+        return super(CategoryModel, self).save(*args, **kwargs)
 
     def save(self, *args, **kwargs):
 
@@ -230,19 +282,20 @@ def create_from_json(json_file):
         for category in data_json.get('categories', []):
             name = category.get('name')
             if name is not None:
-                upsert(CategoryModel, query={ "name": name }, update=category)
+                upsert(CategoryModel, query={"name": name}, update=category)
 
         for dataset_json in data_json.get('datasets', []):
             name = dataset_json.get('name')
             if name:
                 # map category names to ids; create as needed
                 category_ids = []
                 for category in dataset_json.get('categories', []):
-                    category_obj = { "name": category }
+                    category_obj = {"name": category}
+
                     category_model = upsert(CategoryModel, query=category_obj)
                     category_ids.append(category_model.id)
 
                 dataset_json['categories'] = category_ids
                 upsert(DatasetModel, query={ "name": name}, update=dataset_json)
 
-    sys.stdout.flush()
+

app/util/coco_util.py

@@ -1,4 +1,6 @@
 import pycocotools.mask as mask
+import numpy as np
+import cv2
 
 from .query_util import fix_ids
 from ..models import *
@@ -64,6 +66,23 @@ def paperjs_to_coco(image_width, image_height, paperjs):
     return segments, mask.area(rle), mask.toBbox(rle)
 
 
+def get_annotations_iou(annotation_a, annotation_b):
+    """
+    Computes the IOU between two annotation objects
+    """
+    seg_a = list([list(part) for part in annotation_a.segmentation])
+    seg_b = list([list(part) for part in annotation_b.segmentation])
+
+    rles_a = mask.frPyObjects(
+        seg_a, annotation_a.height, annotation_a.width)
+
+    rles_b = mask.frPyObjects(
+        seg_b, annotation_b.height, annotation_b.width)
+
+    ious = mask.iou(rles_a, rles_b, [0])
+    return ious[0][0]
+
+
 def get_image_coco(image):
     """
     Generates coco for an image
@@ -157,6 +176,19 @@ def get_dataset_coco(dataset):
     return coco
 
 
+def decode_seg(mask, segmentation):
+    """
+    Create binary mask from segmentation
+    """
+    pts = [
+        np.array(anno).reshape(-1, 2).round().astype(int)
+        for anno in segmentation
+    ]
+    mask = cv2.fillPoly(mask, pts, 1)
+
+    return mask
+
+
 def _fit(value, max_value, min_value):
 
     if value > max_value: