Add Cuboid2D annotation (#1601)

<!-- Contributing guide:
https://github.com/openvinotoolkit/datumaro/blob/develop/CONTRIBUTING.md
-->

### Summary

This introduces the new annotation type for 3D bounding box. This
annotation is added to the Datumaro format.


<!--
Resolves #111 and #222.
Depends on #1000 (for series of dependent commits).

This PR introduces this capability to make the project better in this
and that.

- Added this feature
- Removed that feature
- Fixed the problem #1234
-->

### How to test
<!-- Describe the testing procedure for reviewers, if changes are
not fully covered by unit tests or manual testing can be complicated.
-->

### Checklist
<!-- Put an 'x' in all the boxes that apply -->
- [x] I have added unit tests to cover my changes.​
- [ ] I have added integration tests to cover my changes.​
- [ ] I have added the description of my changes into
[CHANGELOG](https://github.com/openvinotoolkit/datumaro/blob/develop/CHANGELOG.md).​
- [ ] I have updated the
[documentation](https://github.com/openvinotoolkit/datumaro/tree/develop/docs)
accordingly

### License

- [x] I submit _my code changes_ under the same [MIT
License](https://github.com/openvinotoolkit/datumaro/blob/develop/LICENSE)
that covers the project.
  Feel free to contact the maintainers if that's a concern.
- [ ] I have updated the license header for each file (see an example
below).

```python
# Copyright (C) 2024 Intel Corporation
#
# SPDX-License-Identifier: MIT
```

---------

Signed-off-by: Ilya Trushkin <[email protected]>
Co-authored-by: Wonju Lee <[email protected]>

Author: itrushkin

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### New features
 - Add a new CLI command: datum format
   (<https://github.com/openvinotoolkit/datumaro/pull/1570>)
+- Add a new Cuboid2D annotation type
+  (<https://github.com/openvinotoolkit/datumaro/pull/1601>)
 - Support language dataset for DmTorchDataset
   (<https://github.com/openvinotoolkit/datumaro/pull/1592>)
 

src/datumaro/components/annotation.py

@@ -50,6 +50,7 @@ class AnnotationType(IntEnum):
     feature_vector = 13
     tabular = 14
     rotated_bbox = 15
+    cuboid_2d = 16
 
 
 COORDINATE_ROUNDING_DIGITS = 2
@@ -1363,6 +1364,41 @@ def wrap(item, **kwargs):
         return attr.evolve(item, **d)
 
 
+@attrs(slots=True, init=False, order=False)
+class Cuboid2D(Annotation):
+    """
+    Cuboid2D annotation class. This class represents a 3D bounding box defined by its point coordinates
+    in the following way:
+    [(x1, y1), (x2, y2), (x3, y3), (x4, y4), (x5, y5), (x6, y6), (x7, y7), (x8, y8)].
+
+
+      6---7
+     /|  /|
+    5-+-8 |
+    | 2 + 3
+    |/  |/
+    1---4
+
+    Attributes:
+        _type (AnnotationType): The type of annotation, set to `AnnotationType.bbox`.
+
+    Methods:
+        __init__: Initializes the Cuboid2D with its coordinates.
+        wrap: Creates a new Bbox instance with updated attributes.
+    """
+
+    _type = AnnotationType.cuboid_2d
+    points = field(default=None)
+    label: Optional[int] = field(
+        converter=attr.converters.optional(int), default=None, kw_only=True
+    )
+    z_order: int = field(default=0, validator=default_if_none(int), kw_only=True)
+
+    def __init__(self, _points: Iterable[Tuple[float, float]], *args, **kwargs):
+        kwargs.pop("points", None)  # comes from wrap()
+        self.__attrs_init__(points=_points, *args, **kwargs)
+
+
 @attrs(slots=True, order=False)
 class PointsCategories(Categories):
     """

src/datumaro/components/annotations/matcher.py

@@ -35,6 +35,7 @@
     "ImageAnnotationMatcher",
     "HashKeyMatcher",
     "FeatureVectorMatcher",
+    "Cuboid2DMatcher",
 ]
 
 
@@ -378,3 +379,8 @@ def distance(self, a, b):
         b = Points([p for pt in b.as_polygon() for p in pt])
 
         return OKS(a, b, sigma=self.sigma)
+
+
+@attrs
+class Cuboid2DMatcher(ShapeMatcher):
+    pass

src/datumaro/components/annotations/merger.py

@@ -12,6 +12,7 @@
     AnnotationMatcher,
     BboxMatcher,
     CaptionsMatcher,
+    Cuboid2DMatcher,
     Cuboid3dMatcher,
     FeatureVectorMatcher,
     HashKeyMatcher,
@@ -210,3 +211,8 @@ class TabularMerger(AnnotationMerger, TabularMatcher):
 @attrs
 class RotatedBboxMerger(_ShapeMerger, RotatedBboxMatcher):
     pass
+
+
+@attrs
+class Cuboid2DMerger(_ShapeMerger, Cuboid2DMatcher):
+    pass

src/datumaro/components/merge/intersect_merge.py

@@ -19,6 +19,7 @@
     AnnotationMerger,
     BboxMerger,
     CaptionsMerger,
+    Cuboid2DMerger,
     Cuboid3dMerger,
     EllipseMerger,
     FeatureVectorMerger,
@@ -455,6 +456,8 @@ def _for_type(t, **kwargs):
                 return _make(TabularMerger, **kwargs)
             elif t is AnnotationType.rotated_bbox:
                 return _make(RotatedBboxMerger, **kwargs)
+            elif t is AnnotationType.cuboid_2d:
+                return _make(Cuboid2DMerger, **kwargs)
             else:
                 raise NotImplementedError("Type %s is not supported" % t)
 

src/datumaro/components/visualizer.py

@@ -19,6 +19,7 @@
     AnnotationType,
     Bbox,
     Caption,
+    Cuboid2D,
     Cuboid3d,
     DepthAnnotation,
     Ellipse,
@@ -661,6 +662,39 @@ def _draw_cuboid_3d(
     ) -> None:
         raise NotImplementedError(f"{ann.type} is not implemented yet.")
 
+    def _draw_cuboid_2d(
+        self,
+        ann: Cuboid2D,
+        label_categories: Optional[LabelCategories],
+        fig: Figure,
+        ax: Axes,
+        context: List,
+    ) -> None:
+        import matplotlib.patches as patches
+
+        points = ann.points
+        color = self._get_color(ann)
+        label_text = label_categories[ann.label].name if label_categories is not None else ann.label
+
+        # Define the faces based on vertex indices
+
+        faces = [
+            [points[i] for i in [0, 1, 2, 3]],  # Bottom face
+            [points[i] for i in [4, 5, 6, 7]],  # Top face
+            [points[i] for i in [0, 1, 5, 4]],  # Front face
+            [points[i] for i in [1, 2, 6, 5]],  # Right face
+            [points[i] for i in [2, 3, 7, 6]],  # Back face
+            [points[i] for i in [3, 0, 4, 7]],  # Left face
+        ]
+        ax.text(points[0][0], points[0][1] - self.text_y_offset, label_text, color=color)
+
+        # Draw each face
+        for face in faces:
+            polygon = patches.Polygon(
+                face, fill=False, linewidth=self.bbox_linewidth, edgecolor=color
+            )
+            ax.add_patch(polygon)
+
     def _draw_super_resolution_annotation(
         self,
         ann: SuperResolutionAnnotation,

src/datumaro/plugins/data_formats/datumaro/base.py

@@ -11,6 +11,7 @@
     AnnotationType,
     Bbox,
     Caption,
+    Cuboid2D,
     Cuboid3d,
     Ellipse,
     GroupType,
@@ -378,6 +379,18 @@ def _load_annotations(self, item: Dict):
 
                 elif ann_type == AnnotationType.hash_key:
                     continue
+                elif ann_type == AnnotationType.cuboid_2d:
+                    loaded.append(
+                        Cuboid2D(
+                            list(map(tuple, points)),
+                            label=label_id,
+                            id=ann_id,
+                            attributes=attributes,
+                            group=group,
+                            object_id=object_id,
+                            z_order=z_order,
+                        )
+                    )
                 else:
                     raise NotImplementedError()
             except Exception as e:

src/datumaro/plugins/data_formats/datumaro/exporter.py

@@ -20,6 +20,7 @@
     Annotation,
     Bbox,
     Caption,
+    Cuboid2D,
     Cuboid3d,
     Ellipse,
     HashKey,
@@ -311,6 +312,8 @@ def _gen_item_desc(self, item: DatasetItem, *args, **kwargs) -> Dict:
                 converted_ann = self._convert_ellipse_object(ann)
             elif isinstance(ann, HashKey):
                 continue
+            elif isinstance(ann, Cuboid2D):
+                converted_ann = self._convert_cuboid_2d_object(ann)
             else:
                 raise NotImplementedError()
             annotations.append(converted_ann)
@@ -435,6 +438,18 @@ def _convert_cuboid_3d_object(self, obj):
     def _convert_ellipse_object(self, obj: Ellipse):
         return self._convert_shape_object(obj)
 
+    def _convert_cuboid_2d_object(self, obj: Cuboid2D):
+        converted = self._convert_annotation(obj)
+
+        converted.update(
+            {
+                "label_id": cast(obj.label, int),
+                "points": obj.points,
+                "z_order": obj.z_order,
+            }
+        )
+        return converted
+
 
 class _StreamSubsetWriter(_SubsetWriter):
     def __init__(

src/datumaro/plugins/data_formats/datumaro_binary/mapper/__init__.py

@@ -22,6 +22,7 @@
     "CaptionMapper",
     "Cuboid3dMapper",
     "EllipseMapper",
+    "Cuboid2DMapper",
     # common
     "Mapper",
     "DictMapper",

src/datumaro/plugins/data_formats/datumaro_binary/mapper/annotation.py

@@ -12,6 +12,7 @@
     AnnotationType,
     Bbox,
     Caption,
+    Cuboid2D,
     Cuboid3d,
     Ellipse,
     Label,
@@ -270,6 +271,33 @@ def backward(cls, _bytes: bytes, offset: int = 0) -> Tuple[Ellipse, int]:
         return Ellipse(x, y, x2, y2, **shape_dict), offset
 
 
+class Cuboid2DMapper(AnnotationMapper):
+    ann_type = AnnotationType.cuboid_2d
+
+    @classmethod
+    def forward(cls, ann: Shape) -> bytes:
+        _bytearray = bytearray()
+        _bytearray.extend(struct.pack("<Bqqi", ann.type, ann.id, ann.group, ann.object_id))
+        _bytearray.extend(DictMapper.forward(ann.attributes))
+        _bytearray.extend(
+            struct.pack("<ii", _ShapeMapper.forward_optional_label(ann.label), ann.z_order)
+        )
+        _bytearray.extend(
+            FloatListMapper.forward([coord for point in ann.points for coord in point])
+        )
+        return bytes(_bytearray)
+
+    @classmethod
+    def backward(cls, _bytes: bytes, offset: int = 0) -> Tuple[Ellipse, int]:
+        ann_dict, offset = super().backward_dict(_bytes, offset)
+        label, z_order = struct.unpack_from("<ii", _bytes, offset)
+        offset += 8
+        points, offset = FloatListMapper.backward(_bytes, offset)
+        points = list(zip(points[::2], points[1::2]))
+
+        return Cuboid2D(points, label=label, z_order=z_order, **ann_dict), offset
+
+
 class AnnotationListMapper(Mapper):
     backward_map = {
         AnnotationType.label: LabelMapper.backward,
@@ -281,6 +309,7 @@ class AnnotationListMapper(Mapper):
         AnnotationType.caption: CaptionMapper.backward,
         AnnotationType.cuboid_3d: Cuboid3dMapper.backward,
         AnnotationType.ellipse: EllipseMapper.backward,
+        AnnotationType.cuboid_2d: Cuboid2DMapper.backward,
     }
 
     @classmethod
@@ -307,6 +336,8 @@ def forward(cls, anns: List[Annotation]) -> bytes:
                 _bytearray.extend(Cuboid3dMapper.forward(ann))
             elif isinstance(ann, Ellipse):
                 _bytearray.extend(EllipseMapper.forward(ann))
+            elif isinstance(ann, Cuboid2D):
+                _bytearray.extend(Cuboid2DMapper.forward(ann))
             else:
                 raise NotImplementedError()