Coco converter updates (cvat-ai#864)

Author: zhiltsov-max

datumaro/datumaro/components/config_model.py

@@ -11,7 +11,7 @@
 SOURCE_SCHEMA = _SchemaBuilder() \
     .add('url', str) \
     .add('format', str) \
-    .add('options', str) \
+    .add('options', dict) \
     .build()
 
 class Source(Config):

datumaro/datumaro/components/converters/ms_coco.py

@@ -121,40 +121,96 @@ def save_categories(self, dataset):
             })
 
     def save_annotations(self, item):
-        for ann in item.annotations:
-            if ann.type != AnnotationType.bbox:
+        annotations = item.annotations.copy()
+
+        while len(annotations) != 0:
+            ann = annotations.pop()
+
+            if ann.type == AnnotationType.bbox and ann.label is not None:
+                pass
+            elif ann.type == AnnotationType.polygon and ann.label is not None:
+                pass
+            elif ann.type == AnnotationType.mask and ann.label is not None:
+                pass
+            else:
                 continue
 
-            is_crowd = ann.attributes.get('is_crowd', False)
+            bbox = None
             segmentation = None
-            if ann.group is not None:
+
+            if ann.type == AnnotationType.bbox:
+                is_crowd = ann.attributes.get('is_crowd', False)
+                bbox = ann.get_bbox()
+            elif ann.type == AnnotationType.polygon:
+                is_crowd = ann.attributes.get('is_crowd', False)
+            elif ann.type == AnnotationType.mask:
+                is_crowd = ann.attributes.get('is_crowd', True)
                 if is_crowd:
-                    segmentation = find(item.annotations, lambda x: \
-                        x.group == ann.group and x.type == AnnotationType.mask)
-                    if segmentation is not None:
-                        binary_mask = np.array(segmentation.image, dtype=np.bool)
-                        binary_mask = np.asfortranarray(binary_mask, dtype=np.uint8)
-                        segmentation = mask_utils.encode(binary_mask)
-                        area = mask_utils.area(segmentation)
-                        segmentation = mask_tools.convert_mask_to_rle(binary_mask)
-                else:
-                    segmentation = find(item.annotations, lambda x: \
-                        x.group == ann.group and x.type == AnnotationType.polygon)
-                    if segmentation is not None:
-                        area = ann.area()
-                        segmentation = [segmentation.get_points()]
+                    segmentation = ann
+            area = None
+
+            # If ann in a group, try to find corresponding annotations in
+            # this group, otherwise try to infer them.
+
+            if bbox is None and ann.group is not None:
+                bbox = find(annotations, lambda x: \
+                    x.group == ann.group and \
+                    x.type == AnnotationType.bbox and \
+                    x.label == ann.label)
+                if bbox is not None:
+                    bbox = bbox.get_bbox()
+
+            if is_crowd:
+                # is_crowd=True means there should be a mask
+                if segmentation is None and ann.group is not None:
+                    segmentation = find(annotations, lambda x: \
+                        x.group == ann.group and \
+                        x.type == AnnotationType.mask and \
+                        x.label == ann.label)
+                if segmentation is not None:
+                    binary_mask = np.array(segmentation.image, dtype=np.bool)
+                    binary_mask = np.asfortranarray(binary_mask, dtype=np.uint8)
+                    segmentation = mask_utils.encode(binary_mask)
+                    area = mask_utils.area(segmentation)
+                    segmentation = mask_tools.convert_mask_to_rle(binary_mask)
+            else:
+                # is_crowd=False means there are some polygons
+                polygons = []
+                if ann.type == AnnotationType.polygon:
+                    polygons = [ ann ]
+                if ann.group is not None:
+                    # A single object can consist of several polygons
+                    polygons += [p for p in annotations
+                        if p.group == ann.group and \
+                           p.type == AnnotationType.polygon and \
+                           p.label == ann.label]
+                if polygons:
+                    segmentation = [p.get_points() for p in polygons]
+                    h, w, _ = item.image.shape
+                    rles = mask_utils.frPyObjects(segmentation, h, w)
+                    rle = mask_utils.merge(rles)
+                    area = mask_utils.area(rle)
+
+            if ann.group is not None:
+                # Mark the group as visited to prevent repeats
+                for a in annotations[:]:
+                    if a.group == ann.group:
+                        annotations.remove(a)
+
             if segmentation is None:
                 is_crowd = False
                 segmentation = [ann.get_polygon()]
                 area = ann.area()
+            if bbox is None:
+                bbox = ann.get_bbox()
 
             elem = {
                 'id': self._get_ann_id(ann),
                 'image_id': _cast(item.id, int, 0),
                 'category_id': _cast(ann.label, int, -1) + 1,
                 'segmentation': segmentation,
                 'area': float(area),
-                'bbox': ann.get_bbox(),
+                'bbox': bbox,
                 'iscrowd': int(is_crowd),
             }
             if 'score' in ann.attributes:

datumaro/datumaro/components/extractor.py

@@ -271,6 +271,14 @@ def __init__(self, points=None,
     def get_polygon(self):
         return self.get_points()
 
+    def area(self):
+        import pycocotools.mask as mask_utils
+
+        _, _, w, h = self.get_bbox()
+        rle = mask_utils.frPyObjects([self.get_points()], h, w)
+        area = mask_utils.area(rle)
+        return area
+
 class BboxObject(ShapeObject):
     # pylint: disable=redefined-builtin
     def __init__(self, x=0, y=0, w=0, h=0,

datumaro/datumaro/components/extractors/ms_coco.py

@@ -61,7 +61,7 @@ def __len__(self):
         def categories(self):
             return self._parent.categories()
 
-    def __init__(self, path, task):
+    def __init__(self, path, task, merge_instance_polygons=False):
         super().__init__()
 
         rootpath = path.rsplit(CocoPath.ANNOTATIONS_DIR, maxsplit=1)[0]
@@ -80,6 +80,8 @@ def __init__(self, path, task):
 
         self._load_categories()
 
+        self._merge_instance_polygons = merge_instance_polygons
+
     @staticmethod
     def _make_subset_loader(path):
         # COCO API has an 'unclosed file' warning
@@ -212,20 +214,22 @@ def _parse_annotation(self, ann, ann_type, parsed_annotations,
             segmentation = ann.get('segmentation')
             if segmentation is not None:
                 group = ann_id
+                rle = None
 
                 if isinstance(segmentation, list):
-                    # polygon -- a single object might consist of multiple parts
+                    # polygon - a single object can consist of multiple parts
                     for polygon_points in segmentation:
                         parsed_annotations.append(PolygonObject(
                             points=polygon_points, label=label_id,
-                            group=group
+                            id=ann_id, group=group, attributes=attributes
                         ))
 
-                    # we merge all parts into one mask RLE code
-                    img_h = image_info['height']
-                    img_w = image_info['width']
-                    rles = mask_utils.frPyObjects(segmentation, img_h, img_w)
-                    rle = mask_utils.merge(rles)
+                    if self._merge_instance_polygons:
+                        # merge all parts into a single mask RLE
+                        img_h = image_info['height']
+                        img_w = image_info['width']
+                        rles = mask_utils.frPyObjects(segmentation, img_h, img_w)
+                        rle = mask_utils.merge(rles)
                 elif isinstance(segmentation['counts'], list):
                     # uncompressed RLE
                     img_h, img_w = segmentation['size']
@@ -234,9 +238,10 @@ def _parse_annotation(self, ann, ann_type, parsed_annotations,
                     # compressed RLE
                     rle = segmentation
 
-                parsed_annotations.append(RleMask(rle=rle, label=label_id,
-                    group=group
-                ))
+                if rle is not None:
+                    parsed_annotations.append(RleMask(rle=rle, label=label_id,
+                        id=ann_id, group=group, attributes=attributes
+                    ))
 
             parsed_annotations.append(
                 BboxObject(x, y, w, h, label=label_id,
@@ -277,21 +282,22 @@ def _parse_annotation(self, ann, ann_type, parsed_annotations,
         return parsed_annotations
 
 class CocoImageInfoExtractor(CocoExtractor):
-    def __init__(self, path):
-        super().__init__(path, task=CocoAnnotationType.image_info)
+    def __init__(self, path, **kwargs):
+        super().__init__(path, task=CocoAnnotationType.image_info, **kwargs)
 
 class CocoCaptionsExtractor(CocoExtractor):
-    def __init__(self, path):
-        super().__init__(path, task=CocoAnnotationType.captions)
+    def __init__(self, path, **kwargs):
+        super().__init__(path, task=CocoAnnotationType.captions, **kwargs)
 
 class CocoInstancesExtractor(CocoExtractor):
-    def __init__(self, path):
-        super().__init__(path, task=CocoAnnotationType.instances)
+    def __init__(self, path, **kwargs):
+        super().__init__(path, task=CocoAnnotationType.instances, **kwargs)
 
 class CocoPersonKeypointsExtractor(CocoExtractor):
-    def __init__(self, path):
-        super().__init__(path, task=CocoAnnotationType.person_keypoints)
+    def __init__(self, path, **kwargs):
+        super().__init__(path, task=CocoAnnotationType.person_keypoints,
+            **kwargs)
 
 class CocoLabelsExtractor(CocoExtractor):
-    def __init__(self, path):
-        super().__init__(path, task=CocoAnnotationType.labels)
+    def __init__(self, path, **kwargs):
+        super().__init__(path, task=CocoAnnotationType.labels, **kwargs)

datumaro/datumaro/components/importers/ms_coco.py

@@ -22,7 +22,7 @@ class CocoImporter:
     def __init__(self, task_filter=None):
         self._task_filter = task_filter
 
-    def __call__(self, path):
+    def __call__(self, path, **extra_params):
         from datumaro.components.project import Project # cyclic import
         project = Project()
 
@@ -37,6 +37,7 @@ def __call__(self, path):
                 project.add_source(source_name, {
                     'url': ann_file,
                     'format': self._COCO_EXTRACTORS[ann_type],
+                    'options': extra_params,
                 })
 
         return project