From 136694c04bd6a79469a074726ffa0d53dfc5276b Mon Sep 17 00:00:00 2001 From: Ilya Trushkin Date: Thu, 19 Sep 2024 07:03:47 +0300 Subject: [PATCH] Add Cuboid2D annotation (#1601) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary This introduces the new annotation type for 3D bounding box. This annotation is added to the Datumaro format. ### How to test ### Checklist - [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 Co-authored-by: Wonju Lee --- CHANGELOG.md | 2 + src/datumaro/components/annotation.py | 36 +++++++++++++ .../components/annotations/matcher.py | 6 +++ src/datumaro/components/annotations/merger.py | 6 +++ .../components/merge/intersect_merge.py | 3 ++ src/datumaro/components/visualizer.py | 34 ++++++++++++ .../plugins/data_formats/datumaro/base.py | 13 +++++ .../plugins/data_formats/datumaro/exporter.py | 15 ++++++ .../datumaro_binary/mapper/__init__.py | 1 + .../datumaro_binary/mapper/annotation.py | 31 +++++++++++ tests/unit/data_formats/datumaro/conftest.py | 20 +++++++ tests/unit/operations/test_statistics.py | 52 +++++++++++++++---- 12 files changed, 208 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 161b2d54e9..329161c693 100644 --- a/CHANGELOG.md +++ b/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 () +- Add a new Cuboid2D annotation type + () - Support language dataset for DmTorchDataset () diff --git a/src/datumaro/components/annotation.py b/src/datumaro/components/annotation.py index c68af00c4a..c43599227f 100644 --- a/src/datumaro/components/annotation.py +++ b/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): """ diff --git a/src/datumaro/components/annotations/matcher.py b/src/datumaro/components/annotations/matcher.py index db9322722a..eb7c874cc4 100644 --- a/src/datumaro/components/annotations/matcher.py +++ b/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 diff --git a/src/datumaro/components/annotations/merger.py b/src/datumaro/components/annotations/merger.py index c1c356f81b..8ff7593a61 100644 --- a/src/datumaro/components/annotations/merger.py +++ b/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 diff --git a/src/datumaro/components/merge/intersect_merge.py b/src/datumaro/components/merge/intersect_merge.py index 26677661ea..bb545f950d 100644 --- a/src/datumaro/components/merge/intersect_merge.py +++ b/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) diff --git a/src/datumaro/components/visualizer.py b/src/datumaro/components/visualizer.py index 7030165871..12b4acc05a 100644 --- a/src/datumaro/components/visualizer.py +++ b/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, diff --git a/src/datumaro/plugins/data_formats/datumaro/base.py b/src/datumaro/plugins/data_formats/datumaro/base.py index ee7a8cdc21..a4034269f7 100644 --- a/src/datumaro/plugins/data_formats/datumaro/base.py +++ b/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: diff --git a/src/datumaro/plugins/data_formats/datumaro/exporter.py b/src/datumaro/plugins/data_formats/datumaro/exporter.py index 494492cbe8..d20e019c0e 100644 --- a/src/datumaro/plugins/data_formats/datumaro/exporter.py +++ b/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__( diff --git a/src/datumaro/plugins/data_formats/datumaro_binary/mapper/__init__.py b/src/datumaro/plugins/data_formats/datumaro_binary/mapper/__init__.py index cefedf4cbd..01ee56d60a 100644 --- a/src/datumaro/plugins/data_formats/datumaro_binary/mapper/__init__.py +++ b/src/datumaro/plugins/data_formats/datumaro_binary/mapper/__init__.py @@ -22,6 +22,7 @@ "CaptionMapper", "Cuboid3dMapper", "EllipseMapper", + "Cuboid2DMapper", # common "Mapper", "DictMapper", diff --git a/src/datumaro/plugins/data_formats/datumaro_binary/mapper/annotation.py b/src/datumaro/plugins/data_formats/datumaro_binary/mapper/annotation.py index 4c7269719e..c26658bc64 100644 --- a/src/datumaro/plugins/data_formats/datumaro_binary/mapper/annotation.py +++ b/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(" Tuple[Ellipse, int]: + ann_dict, offset = super().backward_dict(_bytes, offset) + label, z_order = struct.unpack_from(" 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() diff --git a/tests/unit/data_formats/datumaro/conftest.py b/tests/unit/data_formats/datumaro/conftest.py index e600ae957c..693384c40b 100644 --- a/tests/unit/data_formats/datumaro/conftest.py +++ b/tests/unit/data_formats/datumaro/conftest.py @@ -15,6 +15,7 @@ AnnotationType, Bbox, Caption, + Cuboid2D, Cuboid3d, Ellipse, Label, @@ -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( diff --git a/tests/unit/operations/test_statistics.py b/tests/unit/operations/test_statistics.py index bb92c53308..7f28be820a 100644 --- a/tests/unit/operations/test_statistics.py +++ b/tests/unit/operations/test_statistics.py @@ -10,7 +10,16 @@ import numpy as np import pytest -from datumaro.components.annotation import Bbox, Caption, Ellipse, Label, Mask, Points, RotatedBbox +from datumaro.components.annotation import ( + Bbox, + Caption, + Cuboid2D, + Ellipse, + Label, + Mask, + Points, + RotatedBbox, +) from datumaro.components.dataset import Dataset from datumaro.components.dataset_base import DatasetItem from datumaro.components.errors import DatumaroError @@ -232,6 +241,25 @@ def test_stats(self): "tiny": True, }, ), + 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(id=3), @@ -242,7 +270,7 @@ def test_stats(self): expected = { "images count": 4, - "annotations count": 12, + "annotations count": 13, "unannotated images count": 2, "unannotated images": ["3", "2.2"], "annotations by type": { @@ -277,33 +305,34 @@ def test_stats(self): "hash_key": {"count": 0}, "feature_vector": {"count": 0}, "tabular": {"count": 0}, + "cuboid_2d": {"count": 1}, "unknown": {"count": 0}, }, "annotations": { "labels": { - "count": 6, + "count": 7, "distribution": { - "label_0": [1, 1 / 6], + "label_0": [1, 1 / 7], "label_1": [0, 0.0], - "label_2": [3, 3 / 6], - "label_3": [2, 2 / 6], + "label_2": [3, 3 / 7], + "label_3": [3, 3 / 7], }, "attributes": { "x": { - "count": 2, # annotations with no label are skipped + "count": 3, # annotations with no label are skipped "values count": 2, "values present": ["1", "2"], "distribution": { - "1": [1, 1 / 2], - "2": [1, 1 / 2], + "1": [2, 2 / 3], + "2": [1, 1 / 3], }, }, "y": { - "count": 2, # annotations with no label are skipped + "count": 3, # annotations with no label are skipped "values count": 1, "values present": ["2"], "distribution": { - "2": [2, 2 / 2], + "2": [3, 3 / 3], }, }, # must not include "special" attributes like "occluded" @@ -403,6 +432,7 @@ def _get_stats_template(label_names: list): "feature_vector": {"count": 0}, "tabular": {"count": 0}, "rotated_bbox": {"count": 0}, + "cuboid_2d": {"count": 0}, "unknown": {"count": 0}, }, "annotations": {