Skip to content

Commit

Permalink
Add new annotation type RotatedBbox (#1459)
Browse files Browse the repository at this point in the history
<!-- Contributing guide:
https://github.com/openvinotoolkit/datumaro/blob/develop/CONTRIBUTING.md
-->

### Summary

Add a new RotatedBbox annotation type for detection_rotated task.
This provides from_polygon and as_polygon methods for converting set of
coordinates to (cx, cy, w, h, r).

This updates the RoboflowYoloObb behavior as a detection_rotated task.

<!--
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

- [ ] 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
```
  • Loading branch information
wonjuleee authored Apr 22, 2024
1 parent cbdb220 commit e0657d8
Show file tree
Hide file tree
Showing 11 changed files with 205 additions and 39 deletions.
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
97 changes: 97 additions & 0 deletions src/datumaro/components/annotation.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

from __future__ import annotations

import math
from enum import IntEnum
from functools import partial
from itertools import zip_longest
Expand Down Expand Up @@ -48,6 +49,7 @@ class AnnotationType(IntEnum):
hash_key = 12
feature_vector = 13
tabular = 14
rotated_bbox = 15


COORDINATE_ROUNDING_DIGITS = 2
Expand Down Expand Up @@ -845,6 +847,101 @@ def wrap(item, **kwargs):
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):
"""
Expand Down
13 changes: 12 additions & 1 deletion src/datumaro/components/annotations/matcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
7 changes: 7 additions & 0 deletions src/datumaro/components/annotations/merger.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
MaskMatcher,
PointsMatcher,
PolygonMatcher,
RotatedBboxMatcher,
ShapeMatcher,
TabularMatcher,
)
Expand All @@ -29,6 +30,7 @@
"AnnotationMerger",
"LabelMerger",
"BboxMerger",
"RotatedBboxMerger",
"PolygonMerger",
"MaskMerger",
"PointsMerger",
Expand Down Expand Up @@ -203,3 +205,8 @@ class FeatureVectorMerger(AnnotationMerger, FeatureVectorMatcher):
@attrs
class TabularMerger(AnnotationMerger, TabularMatcher):
pass


@attrs
class RotatedBboxMerger(_ShapeMerger, RotatedBboxMatcher):
pass
3 changes: 3 additions & 0 deletions src/datumaro/components/merge/intersect_merge.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
MaskMerger,
PointsMerger,
PolygonMerger,
RotatedBboxMerger,
TabularMerger,
)
from datumaro.components.dataset_base import DatasetItem, IDataset
Expand Down Expand Up @@ -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)

Expand Down
5 changes: 3 additions & 2 deletions src/datumaro/components/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -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: {
Expand All @@ -53,6 +52,7 @@ def __init__(self):
TaskType.segmentation_instance: {
AnnotationType.label,
AnnotationType.bbox,
AnnotationType.rotated_bbox,
AnnotationType.ellipse,
AnnotationType.polygon,
AnnotationType.points,
Expand All @@ -65,6 +65,7 @@ def __init__(self):
TaskType.mixed: {
AnnotationType.label,
AnnotationType.bbox,
AnnotationType.rotated_bbox,
AnnotationType.cuboid_3d,
AnnotationType.ellipse,
AnnotationType.polygon,
Expand Down
6 changes: 3 additions & 3 deletions src/datumaro/plugins/data_formats/roboflow/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
Bbox,
Label,
LabelCategories,
Polygon,
RotatedBbox,
)
from datumaro.components.dataset import DatasetItem
from datumaro.components.dataset_base import SubsetBase
Expand Down Expand Up @@ -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,
Expand Down
34 changes: 25 additions & 9 deletions tests/unit/data_formats/test_roboflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down
46 changes: 23 additions & 23 deletions tests/unit/operations/test_statistics.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -222,6 +222,16 @@ def test_stats(self):
"occluded": False,
},
),
RotatedBbox(
4,
4,
2,
2,
20,
attributes={
"tiny": True,
},
),
],
),
DatasetItem(id=3),
Expand All @@ -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": {
Expand All @@ -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},
Expand Down Expand Up @@ -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": {
Expand Down
Loading

0 comments on commit e0657d8

Please sign in to comment.