Add new annotation type RotatedBbox

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### New features
 - Add task_type property for dataset
   (<https://github.com/openvinotoolkit/datumaro/pull/1422>)
+- Add AnnotationType.rotated_bbox for oriented object detection
+  (<https://github.com/openvinotoolkit/datumaro/pull/1459>)
 
 ### Enhancements
 - Fix ambiguous COCO format detector

src/datumaro/components/annotation.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import math
 from enum import IntEnum
 from functools import partial
 from itertools import zip_longest
@@ -48,6 +49,7 @@
     hash_key = 12
     feature_vector = 13
     tabular = 14
+    rotated_bbox = 15
 
 
 COORDINATE_ROUNDING_DIGITS = 2
@@ -845,6 +847,101 @@
         return attr.evolve(item, **d)
 
 
+@attrs(slots=True, init=False, order=False)
+class RotatedBbox(_Shape):
+    _type = AnnotationType.rotated_bbox
+
+    def __init__(self, cx, cy, w, h, r, *args, **kwargs):
+        kwargs.pop("points", None)  # comes from wrap()
+        self.__attrs_init__([cx, cy, w, h, r], *args, **kwargs)
+
+    @classmethod
+    def from_rectangle(cls, points: List[Tuple[float, float]], *args, **kwargs):
+        assert len(points) == 4, "polygon for a rotated bbox should have only 4 coordinates."
+
+        # Calculate rotation angle
+        rot = math.atan2(points[1][1] - points[0][1], points[1][0] - points[0][0])
+
+        # Calculate the center of the bounding box
+        cx = (points[0][0] + points[2][0]) / 2
+        cy = (points[0][1] + points[2][1]) / 2
+
+        # Calculate the width and height
+        width = math.sqrt((points[1][0] - points[0][0]) ** 2 + (points[1][1] - points[0][1]) ** 2)
+        height = math.sqrt((points[2][0] - points[1][0]) ** 2 + (points[2][1] - points[1][1]) ** 2)
+
+        return cls(cx=cx, cy=cy, w=width, h=height, r=math.degrees(rot), *args, **kwargs)
+
+    @property
+    def cx(self):
+        return self.points[0]
+
+    @property
+    def cy(self):
+        return self.points[1]
+
+    @property
+    def w(self):
+        return self.points[2]
+
+    @property
+    def h(self):
+        return self.points[3]
+
+    @property
+    def r(self):
+        return self.points[4]
+
+    def get_area(self):
+        return self.w * self.h
+
+    def get_bbox(self):
+        polygon = self.as_polygon()
+        xs = [pt[0] for pt in polygon]
+        ys = [pt[1] for pt in polygon]
+
+        return [min(xs), min(ys), max(xs) - min(xs), max(ys) - min(ys)]
+
+    def get_rotated_bbox(self):
+        return [self.cx, self.cy, self.w, self.h, self.r]
+
+    def as_polygon(self) -> List[Tuple[float, float]]:
+        """Convert [center_x, center_y, width, height, rotation] to 4 coordinates for a rotated bounding box."""
+
+        def _rotate_point(x, y, angle):
+            """Rotate a point around another point."""
+            angle_rad = math.radians(angle)
+            cos_theta = math.cos(angle_rad)
+            sin_theta = math.sin(angle_rad)
+            nx = cos_theta * x - sin_theta * y
+            ny = sin_theta * x + cos_theta * y
+            return nx, ny
+
+        # Calculate corner points of the rectangle
+        corners = [
+            (-self.w / 2, -self.h / 2),
+            (self.w / 2, -self.h / 2),
+            (self.w / 2, self.h / 2),
+            (-self.w / 2, self.h / 2),
+        ]
+
+        # Rotate each corner point
+        rotated_corners = [_rotate_point(p[0], p[1], self.r) for p in corners]
+
+        # Translate the rotated points to the original position
+        return [(p[0] + self.cx, p[1] + self.cy) for p in rotated_corners]
+
+    def iou(self, other: _Shape) -> Union[float, Literal[-1]]:
+        from datumaro.util.annotation_util import bbox_iou
+
+        return bbox_iou(self.get_bbox(), other.get_bbox())
+
+    def wrap(item, **kwargs):
+        d = {"x": item.x, "y": item.y, "w": item.w, "h": item.h, "r": item.r}
+        d.update(kwargs)
+        return attr.evolve(item, **d)
+
+
 @attrs(slots=True, order=False)
 class PointsCategories(Categories):
     """

src/datumaro/components/annotations/matcher.py

@@ -9,7 +9,7 @@
 
 from datumaro.components.abstracts import IMergerContext
 from datumaro.components.abstracts.merger import IMatcherContext
-from datumaro.components.annotation import Annotation
+from datumaro.components.annotation import Annotation, Points
 from datumaro.util.annotation_util import (
     OKS,
     approximate_line,
@@ -367,3 +367,14 @@ def match_annotations(self, sources):
 class TabularMatcher(AnnotationMatcher):
     def match_annotations(self, sources):
         raise NotImplementedError()
+
+
+@attrs
+class RotatedBboxMatcher(ShapeMatcher):
+    sigma: Optional[list] = attrib(default=None)
+
+    def distance(self, a, b):
+        a = Points([p for pt in a.as_polygon() for p in pt])
+        b = Points([p for pt in b.as_polygon() for p in pt])
+
+        return OKS(a, b, sigma=self.sigma)

src/datumaro/components/annotations/merger.py

@@ -21,6 +21,7 @@
     MaskMatcher,
     PointsMatcher,
     PolygonMatcher,
+    RotatedBboxMatcher,
     ShapeMatcher,
     TabularMatcher,
 )
@@ -29,6 +30,7 @@
     "AnnotationMerger",
     "LabelMerger",
     "BboxMerger",
+    "RotatedBboxMerger",
     "PolygonMerger",
     "MaskMerger",
     "PointsMerger",
@@ -203,3 +205,8 @@ class FeatureVectorMerger(AnnotationMerger, FeatureVectorMatcher):
 @attrs
 class TabularMerger(AnnotationMerger, TabularMatcher):
     pass
+
+
+@attrs
+class RotatedBboxMerger(_ShapeMerger, RotatedBboxMatcher):
+    pass

src/datumaro/components/merge/intersect_merge.py

@@ -29,6 +29,7 @@
     MaskMerger,
     PointsMerger,
     PolygonMerger,
+    RotatedBboxMerger,
     TabularMerger,
 )
 from datumaro.components.dataset_base import DatasetItem, IDataset
@@ -452,6 +453,8 @@ def _for_type(t, **kwargs):
                 return _make(FeatureVectorMerger, **kwargs)
             elif t is AnnotationType.tabular:
                 return _make(TabularMerger, **kwargs)
+            elif t is AnnotationType.rotated_bbox:
+                return _make(RotatedBboxMerger, **kwargs)
             else:
                 raise NotImplementedError("Type %s is not supported" % t)
 

src/datumaro/components/task.py

@@ -42,8 +42,7 @@ def __init__(self):
                 AnnotationType.points,
             },
             TaskType.detection_rotated: {
-                AnnotationType.label,
-                AnnotationType.polygon,
+                AnnotationType.rotated_bbox,
             },
             TaskType.detection_3d: {AnnotationType.label, AnnotationType.cuboid_3d},
             TaskType.segmentation_semantic: {
@@ -53,6 +52,7 @@ def __init__(self):
             TaskType.segmentation_instance: {
                 AnnotationType.label,
                 AnnotationType.bbox,
+                AnnotationType.rotated_bbox,
                 AnnotationType.ellipse,
                 AnnotationType.polygon,
                 AnnotationType.points,
@@ -65,6 +65,7 @@ def __init__(self):
             TaskType.mixed: {
                 AnnotationType.label,
                 AnnotationType.bbox,
+                AnnotationType.rotated_bbox,
                 AnnotationType.cuboid_3d,
                 AnnotationType.ellipse,
                 AnnotationType.polygon,

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

@@ -16,7 +16,7 @@
     Bbox,
     Label,
     LabelCategories,
-    Polygon,
+    RotatedBbox,
 )
 from datumaro.components.dataset import DatasetItem
 from datumaro.components.dataset_base import SubsetBase
@@ -171,8 +171,8 @@ def _parse_annotations(
             x4 = self._parse_field(x4, float, "x4")
             y4 = self._parse_field(y4, float, "y4")
             annotations.append(
-                Polygon(
-                    points=[x1, y1, x2, y2, x3, y3, x4, y4],
+                RotatedBbox.from_rectangle(
+                    points=[(x1, y1), (x2, y2), (x3, y3), (x4, y4)],
                     label=label_id,
                     id=idx,
                     group=idx,

tests/unit/data_formats/test_roboflow.py

@@ -7,7 +7,7 @@
 import numpy as np
 import pytest
 
-from datumaro.components.annotation import Bbox, Label, Polygon
+from datumaro.components.annotation import Bbox, Label, RotatedBbox
 from datumaro.components.dataset import Dataset
 from datumaro.components.dataset_base import DatasetItem
 from datumaro.components.environment import DEFAULT_ENVIRONMENT
@@ -174,8 +174,12 @@ def fxt_yolo_obb_dataset():
                 subset="train",
                 media=Image.from_numpy(data=np.ones((5, 10, 3))),
                 annotations=[
-                    Polygon(
-                        points=[0, 0, 0, 2, 2, 2, 2, 0],
+                    RotatedBbox(
+                        1,
+                        1,
+                        2,
+                        2,
+                        90,
                         label=0,
                         group=0,
                         id=0,
@@ -187,8 +191,12 @@ def fxt_yolo_obb_dataset():
                 subset="train",
                 media=Image.from_numpy(data=np.ones((5, 10, 3))),
                 annotations=[
-                    Polygon(
-                        points=[1, 1, 1, 5, 5, 5, 5, 1],
+                    RotatedBbox(
+                        3,
+                        3,
+                        4,
+                        4,
+                        90,
                         label=1,
                         group=0,
                         id=0,
@@ -200,14 +208,22 @@ def fxt_yolo_obb_dataset():
                 subset="val",
                 media=Image.from_numpy(data=np.ones((5, 10, 3))),
                 annotations=[
-                    Polygon(
-                        points=[0, 0, 0, 1, 1, 1, 1, 0],
+                    RotatedBbox(
+                        0.5,
+                        0.5,
+                        1,
+                        1,
+                        90,
                         label=0,
                         group=0,
                         id=0,
                     ),
-                    Polygon(
-                        points=[1, 2, 1, 5, 10, 5, 10, 2],
+                    RotatedBbox(
+                        5.5,
+                        3.5,
+                        3,
+                        9,
+                        90,
                         label=1,
                         group=1,
                         id=1,

tests/unit/operations/test_statistics.py

@@ -10,7 +10,7 @@
 import numpy as np
 import pytest
 
-from datumaro.components.annotation import Bbox, Caption, Ellipse, Label, Mask, Points
+from datumaro.components.annotation import Bbox, Caption, Ellipse, Label, Mask, Points, RotatedBbox
 from datumaro.components.dataset import Dataset
 from datumaro.components.dataset_base import DatasetItem
 from datumaro.components.errors import DatumaroError
@@ -222,6 +222,16 @@ def test_stats(self):
                                 "occluded": False,
                             },
                         ),
+                        RotatedBbox(
+                            4,
+                            4,
+                            2,
+                            2,
+                            20,
+                            attributes={
+                                "tiny": True,
+                            },
+                        ),
                     ],
                 ),
                 DatasetItem(id=3),
@@ -232,7 +242,7 @@ def test_stats(self):
 
         expected = {
             "images count": 4,
-            "annotations count": 11,
+            "annotations count": 12,
             "unannotated images count": 2,
             "unannotated images": ["3", "2.2"],
             "annotations by type": {
@@ -257,6 +267,9 @@ def test_stats(self):
                 "caption": {
                     "count": 2,
                 },
+                "rotated_bbox": {
+                    "count": 1,
+                },
                 "cuboid_3d": {"count": 0},
                 "super_resolution_annotation": {"count": 0},
                 "depth_annotation": {"count": 0},
@@ -375,34 +388,21 @@ def _get_stats_template(label_names: list):
             "unannotated images count": 0,
             "unannotated images": [],
             "annotations by type": {
-                "label": {
-                    "count": 0,
-                },
-                "polygon": {
-                    "count": 0,
-                },
-                "polyline": {
-                    "count": 0,
-                },
-                "bbox": {
-                    "count": 0,
-                },
-                "mask": {
-                    "count": 0,
-                },
-                "points": {
-                    "count": 0,
-                },
-                "caption": {
-                    "count": 0,
-                },
+                "label": {"count": 0},
+                "polygon": {"count": 0},
+                "polyline": {"count": 0},
+                "bbox": {"count": 0},
+                "mask": {"count": 0},
+                "points": {"count": 0},
+                "caption": {"count": 0},
                 "cuboid_3d": {"count": 0},
                 "super_resolution_annotation": {"count": 0},
                 "depth_annotation": {"count": 0},
                 "ellipse": {"count": 0},
                 "hash_key": {"count": 0},
                 "feature_vector": {"count": 0},
                 "tabular": {"count": 0},
+                "rotated_bbox": {"count": 0},
                 "unknown": {"count": 0},
             },
             "annotations": {