Exports only COCO properties

Author: jsbroks

backend/database/__init__.py

@@ -40,9 +40,9 @@ def upsert(model, query=None, update=None):
     return new_model
 
 
-def fix_ids(objs):
-    objects_list = json.loads(objs.to_json().replace('\"_id\"', '\"id\"'))
-    return objects_list
+def fix_ids(q):
+    json_obj = json.loads(q.to_json().replace('\"_id\"', '\"id\"'))
+    return json_obj
 
 
 def create_from_json(json_file):

backend/database/annotations.py

@@ -10,6 +10,10 @@
 
 class AnnotationModel(DynamicDocument):
 
+    COCO_PROPERTIES = ["id", "image_id", "category_id", "segmentation", \
+                       "iscrowd", "color", "area", "bbox", "metadata", \
+                       "keypoints"]
+
     id = SequenceField(primary_key=True)
     image_id = IntField(required=True)
     category_id = IntField(required=True)
@@ -41,13 +45,15 @@ class AnnotationModel(DynamicDocument):
     def __init__(self, image_id=None, **data):
 
         from .images import ImageModel
-        image = ImageModel.objects(id=image_id).first()
 
-        if image is not None:
-            data['image_id'] = image_id
-            data['width'] = image.width
-            data['height'] = image.height
-            data['dataset_id'] = image.dataset_id
+        if image_id is not None:
+            image = ImageModel.objects(id=image_id).first()
+
+            if image is not None:
+                data['image_id'] = image_id
+                data['width'] = image.width
+                data['height'] = image.height
+                data['dataset_id'] = image.dataset_id
 
         super(AnnotationModel, self).__init__(**data)
 

backend/database/categories.py

@@ -7,6 +7,9 @@
 
 class CategoryModel(DynamicDocument):
 
+    COCO_PROPERTIES = ["id", "name", "supercategory", "color", "metadata",\
+                       "keypoint_edges", "keypoint_labels"]
+
     id = SequenceField(primary_key=True)
     name = StringField(required=True, unique_with=['creator'])
     supercategory = StringField(default='')

backend/database/images.py

@@ -10,7 +10,10 @@
 from .annotations import AnnotationModel
 
 class ImageModel(DynamicDocument):
-    
+
+    COCO_PROPERTIES = ["id", "width", "height", "file_name", "path", "license",\
+                       "flickr_url", "coco_url", "date_captured", "dataset_id"]
+
     # -- Contants
     THUMBNAIL_DIRECTORY = '.thumbnail'
     PATTERN = (".gif", ".png", ".jpg", ".jpeg", ".bmp")
@@ -37,6 +40,7 @@ class ImageModel(DynamicDocument):
     thumbnail_url = StringField()
     image_url = StringField()
     coco_url = StringField()
+    date_captured = DateTimeField()
 
     metadata = DictField()
     license = IntField()

backend/webserver/api/images.py

@@ -204,5 +204,5 @@ def get(self, image_id):
         if not current_user.can_download(image):
             return {"message": "You do not have permission to download the images's annotations"}, 403
 
-        return coco_util.get_image_coco(image)
+        return coco_util.get_image_coco(image_id)
 

backend/webserver/util/coco_util.py

@@ -1,7 +1,7 @@
 import pycocotools.mask as mask
 
-from .query_util import fix_ids
 from database import (
+    fix_ids,
     ImageModel,
     DatasetModel,
     CategoryModel,
@@ -102,52 +102,52 @@ def get_annotations_iou(annotation_a, annotation_b):
     return ious[0][0]
 
 
-def get_image_coco(image):
+def get_image_coco(image_id):
     """
     Generates coco for an image
 
     :param image: ImageModel
     :return: Coco in dictionary format
     """
-    dataset = DatasetModel.objects(id=image.dataset_id).first()
-    image = fix_ids(image)
+    image = ImageModel.objects(id=image_id)\
+        .only(*ImageModel.COCO_PROPERTIES)
+    
+    image = fix_ids(image)[0]
+    dataset = DatasetModel.objects(id=image.get('dataset_id')).first()
 
-    bulk_categories = CategoryModel.objects(deleted=False) \
-        .exclude('deleted_date').in_bulk(dataset.categories).items()
+    bulk_categories = CategoryModel.objects(id__in=dataset.categories, deleted=False) \
+        .only(*CategoryModel.COCO_PROPERTIES)
 
+    print(bulk_categories)
+
+    db_annotations = AnnotationModel.objects(deleted=False, image_id=image_id)
     categories = []
     annotations = []
 
-    for category in bulk_categories:
-        category = category[1]
-        category_annotations = AnnotationModel.objects(
-            deleted=False, category_id=category.id, image_id=image.get('id')
-        ).exclude('paper_object', 'deleted_date').all()
+    for category in fix_ids(bulk_categories):
+
+        category_annotations = db_annotations\
+            .filter(category_id=category.get('id'))\
+            .only(*AnnotationModel.COCO_PROPERTIES)
 
-        if len(category_annotations) == 0:
+        if category_annotations.count() == 0:
             continue
 
+        category_annotations = fix_ids(category_annotations)
         for annotation in category_annotations:
-            annotation = fix_ids(annotation)
 
             has_segmentation = len(annotation.get('segmentation', [])) > 0
             has_keypoints = len(annotation.get('keypoints', [])) > 0
 
             if has_segmentation or has_keypoints:
-                del annotation['deleted']
 
-                if not has_keypoints:
-                    del annotation['keypoints']
-                else:
+                if has_keypoints:
                     arr = np.array(annotation.get('keypoints', []))
                     arr = arr[2::3]
                     annotation['num_keypoints'] = len(arr[arr > 0])
-
+                
                 annotations.append(annotation)
 
-        category = fix_ids(category)
-        del category['deleted']
-
         if len(category.get('keypoint_labels')) > 0:
             category['keypoints'] = category.pop('keypoint_labels')
             category['skeleton'] = category.pop('keypoint_edges')
@@ -157,8 +157,6 @@ def get_image_coco(image):
 
         categories.append(category)
 
-    del image['deleted']
-
     coco = {
         "images": [image],
         "categories": categories,

backend/workers/tasks/data.py

@@ -27,34 +27,34 @@ def export_annotations(task_id, dataset_id):
 
     task.info("Beginning Export (COCO Format)")
 
-    categories = CategoryModel.objects(deleted=False) \
-        .exclude('deleted_date').in_bulk(dataset.categories).items()
+    db_categories = CategoryModel.objects(id__in=dataset.categories, deleted=False) \
+        .only(*CategoryModel.COCO_PROPERTIES)
+    db_images = ImageModel.objects(deleted=False, annotated=True, dataset_id=dataset.id)\
+        .only(*ImageModel.COCO_PROPERTIES)
+    db_annotations = AnnotationModel.objects(deleted=False)
 
-    total_items = len(dataset.categories)
+    total_items = db_categories.count()
     dataset = fix_ids(dataset)
 
-    images = ImageModel.objects(deleted=False, annotated=True, dataset_id=dataset.get('id')).exclude('deleted_date')
-    all_annotations = AnnotationModel.objects(deleted=False).exclude('deleted_date', 'paper_object')
-
     coco = {
         'images': [],
         'categories': [],
         'annotations': []
     }
 
-    total_items += images.count()
+    total_items += db_images.count()
     progress = 0
-    for category in categories:
-        category = fix_ids(category[1])
 
-        del category['deleted']
+    # iterate though all categoires and upsert
+    for category in fix_ids(db_categories):
+
         if len(category.get('keypoint_labels', [])) > 0:
             category['keypoints'] = category.pop('keypoint_labels', [])
             category['skeleton'] = category.pop('keypoint_edges', [])
         else:
             if 'keypoint_edges' in category:
                 del category['keypoint_edges']
-            if 'keypoint_labels' in categories:
+            if 'keypoint_labels' in category:
                 del category['keypoint_labels']
 
         task.info(f"Adding category: {category.get('name')}")
@@ -63,22 +63,23 @@ def export_annotations(task_id, dataset_id):
         progress += 1
         task.set_progress((progress/total_items)*100, socket=socket)
 
-    total_annotations = 0
-    total_images = 0
-    for image in images:
-        annotations = all_annotations.filter(image_id=image.id)
-        if annotations.count() == 0:
-            continue
+    total_annotations = db_annotations.count()
+    total_images = db_images.count()
+    for image in fix_ids(db_images):
+
+        progress += 1
+        task.set_progress((progress/total_items)*100, socket=socket)  
 
-        annotations = fix_ids(annotations.all())
+        annotations = db_annotations.filter(image_id=image.get('id'))\
+            .only(*AnnotationModel.COCO_PROPERTIES)
+        annotations = fix_ids(annotations)
         num_annotations = 0
         for annotation in annotations:
 
             has_keypoints = len(annotation.get('keypoints', [])) > 0
             has_segmentation = len(annotation.get('segmentation', [])) > 0
 
             if has_keypoints or has_segmentation:
-                del annotation['deleted']
 
                 if not has_keypoints:
                     if 'keypoints' in annotation:
@@ -91,16 +92,8 @@ def export_annotations(task_id, dataset_id):
                 num_annotations += 1
                 coco.get('annotations').append(annotation)
 
-        task.info(f'Exporting {num_annotations} annotations for image {image.id}')
-        total_annotations += num_annotations
-
-        image = fix_ids(image)
-        del image['deleted']
+        task.info(f"Exporting {num_annotations} annotations for image {image.get('id')}")
         coco.get('images').append(image)
-        total_images += 1
-        
-        progress += 1
-        task.set_progress((progress/total_items)*100, socket=socket)  
 
     file_path = dataset.get('directory') + '/coco.json'
     with open(file_path, 'w') as fp: