Merge pull request cvat-ai#66 from openvinotoolkit/ay/transform-labels

Add function to transform labels

Author: yasakova-anastasia

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 ### Added
 - `WiderFace` dataset format (<https://github.com/openvinotoolkit/datumaro/pull/65>)
+- Function to transform annotations to labels (<https://github.com/openvinotoolkit/datumaro/pull/66>)
 
 ### Changed
 -

datumaro/plugins/transforms.py

@@ -12,7 +12,7 @@
 import pycocotools.mask as mask_utils
 
 from datumaro.components.extractor import (Transform, AnnotationType,
-    RleMask, Polygon, Bbox, DEFAULT_SUBSET_NAME,
+    RleMask, Polygon, Bbox, Label, DEFAULT_SUBSET_NAME,
     LabelCategories, MaskCategories, PointsCategories
 )
 from datumaro.components.cli_plugin import CliPlugin
@@ -541,4 +541,19 @@ def transform_item(self, item):
                     annotations.append(ann.wrap(label=conv_label))
             else:
                 annotations.append(ann.wrap())
+        return item.wrap(annotations=annotations)
+
+class AnnsToLabels(Transform, CliPlugin):
+    """
+    Collects all labels from annotations (of all types) and
+    transforms them into a set of annotations of type Label
+    """
+
+    def transform_item(self, item):
+        labels = set(p.label for p in item.annotations
+            if getattr(p, 'label') != None)
+        annotations = []
+        for label in labels:
+            annotations.append(Label(label=label))
+
         return item.wrap(annotations=annotations)

tests/test_transforms.py

@@ -386,3 +386,29 @@ def test_remap_labels_delete_unspecified(self):
             mapping={}, default='delete')
 
         compare_datasets(self, target_dataset, actual)
+
+    def test_transform_labels(self):
+        src_dataset = Dataset.from_iterable([
+            DatasetItem(id=1, annotations=[
+                Label(1),
+                Bbox(1, 2, 3, 4, label=2),
+                Bbox(1, 3, 3, 3),
+                Mask(image=np.array([1]), label=3),
+                Polygon([1, 1, 2, 2, 3, 4], label=4),
+                PolyLine([1, 3, 4, 2, 5, 6], label=5)
+            ])
+        ], categories=['label%s' % i for i in range(6)])
+
+        dst_dataset = Dataset.from_iterable([
+            DatasetItem(id=1, annotations=[
+                Label(1),
+                Label(2),
+                Label(3),
+                Label(4),
+                Label(5)
+            ]),
+        ], categories=['label%s' % i for i in range(6)])
+
+        actual = transforms.AnnsToLabels(src_dataset)
+
+        compare_datasets(self, dst_dataset, actual)