Allow attributes in COCO format (#1782)

* allow attributes in coco

* update changelog

* update cli

* update coco cli

CHANGELOG.md

@@ -34,6 +34,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Implemented import and export of annotations with relative image paths (<https://github.com/opencv/cvat/pull/1463>)
 - Using only single click to start editing or remove a point (<https://github.com/opencv/cvat/pull/1571>)
 - Added support for attributes in VOC XML format (https://github.com/opencv/cvat/pull/1792)
+- Added annotation attributes in COCO format (https://github.com/opencv/cvat/pull/1782)
 - Colorized object items in the side panel (<https://github.com/opencv/cvat/pull/1753>)
 
 ### Deprecated

datumaro/datumaro/plugins/coco_format/converter.py

@@ -17,7 +17,7 @@ from datumaro.components.extractor import (DEFAULT_SUBSET_NAME,
     AnnotationType, Points
 )
 from datumaro.components.cli_plugin import CliPlugin
-from datumaro.util import find, cast
+from datumaro.util import find, cast, str_to_bool
 from datumaro.util.image import save_image
 import datumaro.util.mask_tools as mask_tools
 import datumaro.util.annotation_tools as anno_tools
@@ -110,6 +110,12 @@ class _TaskConverter:
             self._min_ann_id = max(ann_id, self._min_ann_id)
         return ann_id
 
+    @staticmethod
+    def _convert_attributes(ann):
+        return { k: v for k, v in ann.attributes.items()
+            if k not in {'is_crowd', 'score'}
+        }
+
 class _ImageInfoConverter(_TaskConverter):
     def is_empty(self):
         return len(self._data['images']) == 0
@@ -141,6 +147,8 @@ class _CaptionsConverter(_TaskConverter):
                 except Exception as e:
                     log.warning("Item '%s', ann #%s: failed to convert "
                         "attribute 'score': %e" % (item.id, ann_idx, e))
+            if self._context._allow_attributes:
+                elem['attributes'] = self._convert_attributes(ann)
 
             self.annotations.append(elem)
 
@@ -312,6 +320,8 @@ class _InstancesConverter(_TaskConverter):
             except Exception as e:
                 log.warning("Item '%s': failed to convert attribute "
                     "'score': %e" % (item.id, e))
+        if self._context._allow_attributes:
+                elem['attributes'] = self._convert_attributes(ann)
 
         return elem
 
@@ -428,6 +438,8 @@ class _LabelsConverter(_TaskConverter):
                 except Exception as e:
                     log.warning("Item '%s': failed to convert attribute "
                         "'score': %e" % (item.id, e))
+            if self._context._allow_attributes:
+                elem['attributes'] = self._convert_attributes(ann)
 
             self.annotations.append(elem)
 
@@ -442,7 +454,7 @@ class _Converter:
 
     def __init__(self, extractor, save_dir,
             tasks=None, save_images=False, segmentation_mode=None,
-            crop_covered=False):
+            crop_covered=False, allow_attributes=True):
         assert tasks is None or isinstance(tasks, (CocoTask, list, str))
         if tasks is None:
             tasks = list(self._TASK_CONVERTER)
@@ -473,6 +485,7 @@ class _Converter:
         self._segmentation_mode = segmentation_mode
 
         self._crop_covered = crop_covered
+        self._allow_attributes = allow_attributes
 
         self._image_ids = {}
 
@@ -549,25 +562,29 @@ class CocoConverter(Converter, CliPlugin):
 
     @classmethod
     def build_cmdline_parser(cls, **kwargs):
+        kwargs['description'] = """
+            Segmentation save modes:|n
+            - '{sm.guess.name}': guess the mode for each instance,|n
+            |s|suse 'is_crowd' attribute as a hint|n
+            - '{sm.polygons.name}': save polygons,|n
+            |s|smerge and convert masks, prefer polygons|n
+            - '{sm.mask.name}': save masks,|n
+            |s|smerge and convert polygons, prefer masks
+            """.format(sm=SegmentationMode)
         parser = super().build_cmdline_parser(**kwargs)
+
         parser.add_argument('--save-images', action='store_true',
             help="Save images (default: %(default)s)")
         parser.add_argument('--segmentation-mode',
             choices=[m.name for m in SegmentationMode],
             default=SegmentationMode.guess.name,
-            help="""
-                Save mode for instance segmentation:|n
-                - '{sm.guess.name}': guess the mode for each instance,|n
-                |s|suse 'is_crowd' attribute as hint|n
-                - '{sm.polygons.name}': save polygons,|n
-                |s|smerge and convert masks, prefer polygons|n
-                - '{sm.mask.name}': save masks,|n
-                |s|smerge and convert polygons, prefer masks|n
-                Default: %(default)s.
-                """.format(sm=SegmentationMode))
+            help="Save mode for instance segmentation (default: %(default)s)")
         parser.add_argument('--crop-covered', action='store_true',
             help="Crop covered segments so that background objects' "
                 "segmentation was more accurate (default: %(default)s)")
+        parser.add_argument('--allow-attributes',
+            type=str_to_bool, default=True,
+            help="Allow export of attributes (default: %(default)s)")
         parser.add_argument('--tasks', type=cls._split_tasks_string,
             default=None,
             help="COCO task filter, comma-separated list of {%s} "
@@ -576,7 +593,7 @@ class CocoConverter(Converter, CliPlugin):
 
     def __init__(self,
             tasks=None, save_images=False, segmentation_mode=None,
-            crop_covered=False):
+            crop_covered=False, allow_attributes=True):
         super().__init__()
 
         self._options = {
@@ -584,6 +601,7 @@ class CocoConverter(Converter, CliPlugin):
             'save_images': save_images,
             'segmentation_mode': segmentation_mode,
             'crop_covered': crop_covered,
+            'allow_attributes': allow_attributes,
         }
 
     def __call__(self, extractor, save_dir):

datumaro/datumaro/plugins/coco_format/extractor.py

@@ -145,6 +145,12 @@ class _CocoExtractor(SourceExtractor):
         ann_id = ann.get('id')
 
         attributes = {}
+        if 'attributes' in ann:
+            try:
+                attributes.update(ann['attributes'])
+            except Exception as e:
+                log.debug("item #%s: failed to read annotation attributes: %s",
+                    image_info['id'], e)
         if 'score' in ann:
             attributes['score'] = ann['score']
 

datumaro/tests/test_coco_format.py

@@ -574,3 +574,23 @@ class CocoConverterTest(TestCase):
         with TestDir() as test_dir:
             self._test_save_and_load(TestExtractor(),
                 CocoConverter(tasks='image_info', save_images=True), test_dir)
+
+    def test_annotation_attributes(self):
+        class TestExtractor(Extractor):
+            def __iter__(self):
+                return iter([
+                    DatasetItem(id=1, image=np.ones((4, 2, 3)), annotations=[
+                        Polygon([0, 0, 4, 0, 4, 4], label=5, group=1, id=1,
+                            attributes={'is_crowd': False, 'x': 5, 'y': 'abc'}),
+                    ], attributes={'id': 1})
+                ])
+
+            def categories(self):
+                label_categories = LabelCategories()
+                for i in range(10):
+                    label_categories.add(str(i))
+                return { AnnotationType.label: label_categories, }
+
+        with TestDir() as test_dir:
+            self._test_save_and_load(TestExtractor(),
+                CocoConverter(), test_dir)