allow attributes in coco

datumaro/datumaro/plugins/coco_format/converter.py

@@ -17,7 +17,7 @@
     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 @@ def _get_ann_id(self, annotation):
             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 @@ def save_annotations(self, item):
                 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 @@ def convert_instance(self, instance, item):
             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 @@ def save_annotations(self, item):
                 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 @@ def __init__(self, extractor, save_dir,
         self._segmentation_mode = segmentation_mode
 
         self._crop_covered = crop_covered
+        self._allow_attributes = allow_attributes
 
         self._image_ids = {}
 
@@ -568,6 +581,9 @@ def build_cmdline_parser(cls, **kwargs):
         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} "

datumaro/datumaro/plugins/coco_format/extractor.py

@@ -145,6 +145,12 @@ def _load_annotations(self, ann, image_info=None):
         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 @@ def __iter__(self):
         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)