diff --git a/CHANGELOG.md b/CHANGELOG.md index 36f0bc91656..14907f1b420 100644 --- a/CHANGELOG.md +++ b/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 () - Using only single click to start editing or remove a point () - 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 () ### Deprecated diff --git a/datumaro/datumaro/plugins/coco_format/converter.py b/datumaro/datumaro/plugins/coco_format/converter.py index d5247bdbd76..0e7dd3a444c 100644 --- a/datumaro/datumaro/plugins/coco_format/converter.py +++ b/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 = {} @@ -549,25 +562,29 @@ def _split_tasks_string(s): @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 @@ def build_cmdline_parser(cls, **kwargs): 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 @@ def __init__(self, 'save_images': save_images, 'segmentation_mode': segmentation_mode, 'crop_covered': crop_covered, + 'allow_attributes': allow_attributes, } def __call__(self, extractor, save_dir): diff --git a/datumaro/datumaro/plugins/coco_format/extractor.py b/datumaro/datumaro/plugins/coco_format/extractor.py index 8ba0d87d2b3..8bb6e464708 100644 --- a/datumaro/datumaro/plugins/coco_format/extractor.py +++ b/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'] diff --git a/datumaro/tests/test_coco_format.py b/datumaro/tests/test_coco_format.py index 7a31180bea4..08c75d0f2fd 100644 --- a/datumaro/tests/test_coco_format.py +++ b/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) \ No newline at end of file