Skip to content

Commit

Permalink
Add Cuboid2D annotation (openvinotoolkit#1601)
Browse files Browse the repository at this point in the history
<!-- 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 openvinotoolkit#111 and openvinotoolkit#222.
Depends on openvinotoolkit#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 openvinotoolkit#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 <ilya.trushkin@intel.com>
Co-authored-by: Wonju Lee <wonju.lee@intel.com>
  • Loading branch information
itrushkin and wonjuleee authored Sep 19, 2024
1 parent dc63222 commit 136694c
Show file tree
Hide file tree
Showing 12 changed files with 208 additions and 11 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 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>)

Expand Down
36 changes: 36 additions & 0 deletions src/datumaro/components/annotation.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ class AnnotationType(IntEnum):
feature_vector = 13
tabular = 14
rotated_bbox = 15
cuboid_2d = 16


COORDINATE_ROUNDING_DIGITS = 2
Expand Down Expand Up @@ -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):
"""
Expand Down
6 changes: 6 additions & 0 deletions src/datumaro/components/annotations/matcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
"ImageAnnotationMatcher",
"HashKeyMatcher",
"FeatureVectorMatcher",
"Cuboid2DMatcher",
]


Expand Down Expand Up @@ -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
6 changes: 6 additions & 0 deletions src/datumaro/components/annotations/merger.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
AnnotationMatcher,
BboxMatcher,
CaptionsMatcher,
Cuboid2DMatcher,
Cuboid3dMatcher,
FeatureVectorMatcher,
HashKeyMatcher,
Expand Down Expand Up @@ -210,3 +211,8 @@ class TabularMerger(AnnotationMerger, TabularMatcher):
@attrs
class RotatedBboxMerger(_ShapeMerger, RotatedBboxMatcher):
pass


@attrs
class Cuboid2DMerger(_ShapeMerger, Cuboid2DMatcher):
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 @@ -19,6 +19,7 @@
AnnotationMerger,
BboxMerger,
CaptionsMerger,
Cuboid2DMerger,
Cuboid3dMerger,
EllipseMerger,
FeatureVectorMerger,
Expand Down Expand Up @@ -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)

Expand Down
34 changes: 34 additions & 0 deletions src/datumaro/components/visualizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
AnnotationType,
Bbox,
Caption,
Cuboid2D,
Cuboid3d,
DepthAnnotation,
Ellipse,
Expand Down Expand Up @@ -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,
Expand Down
13 changes: 13 additions & 0 deletions src/datumaro/plugins/data_formats/datumaro/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
AnnotationType,
Bbox,
Caption,
Cuboid2D,
Cuboid3d,
Ellipse,
GroupType,
Expand Down Expand Up @@ -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:
Expand Down
15 changes: 15 additions & 0 deletions src/datumaro/plugins/data_formats/datumaro/exporter.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
Annotation,
Bbox,
Caption,
Cuboid2D,
Cuboid3d,
Ellipse,
HashKey,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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__(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
"CaptionMapper",
"Cuboid3dMapper",
"EllipseMapper",
"Cuboid2DMapper",
# common
"Mapper",
"DictMapper",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
AnnotationType,
Bbox,
Caption,
Cuboid2D,
Cuboid3d,
Ellipse,
Label,
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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()

Expand Down
20 changes: 20 additions & 0 deletions tests/unit/data_formats/datumaro/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
AnnotationType,
Bbox,
Caption,
Cuboid2D,
Cuboid3d,
Ellipse,
Label,
Expand Down Expand Up @@ -122,6 +123,25 @@ def fxt_test_datumaro_format_dataset():
"y": "2",
},
),
Cuboid2D(
[
(1, 1),
(3, 1),
(3, 3),
(1, 3),
(1.5, 1.5),
(3.5, 1.5),
(3.5, 3.5),
(1.5, 3.5),
],
label=3,
id=5,
z_order=2,
attributes={
"x": 1,
"y": "2",
},
),
],
),
DatasetItem(
Expand Down
Loading

0 comments on commit 136694c

Please sign in to comment.