Change sematic segmentation to consider bbox only annotations. (#3996)

* segmentation consider bbox only annotations

* add unit test

* add unit test

* update fixture

* use name attribute

* revert tox file

* update for 2.2.0rc4

---------

Co-authored-by: Yunchu Lee <[email protected]>

Author: eunwoosh

CHANGELOG.md

@@ -55,6 +55,8 @@ All notable changes to this project will be documented in this file.
   (<https://github.com/openvinotoolkit/training_extensions/pull/3964>)
 - Add type checker in converter for callable functions (optimizer, scheduler)
   (<https://github.com/openvinotoolkit/training_extensions/pull/3968>)
+- Change sematic segmentation to consider bbox only annotations
+  (<https://github.com/openvinotoolkit/training_extensions/pull/3996>)
 
 ### Bug fixes
 

README.md

@@ -196,6 +196,7 @@ In addition to the examples above, please refer to the documentation for tutoria
 - Include Geti arrow dataset subset names
 - Include full image with anno in case there's no tile in tile dataset
 - Add type checker in converter for callable functions (optimizer, scheduler)
+- Change sematic segmentation to consider bbox only annotations
 
 ### Bug fixes
 

docs/source/guide/release_notes/index.rst

@@ -37,6 +37,7 @@ Enhancements
 - Include Geti arrow dataset subset names
 - Include full image with anno in case there's no tile in tile dataset
 - Add type checker in converter for callable functions (optimizer, scheduler)
+- Change sematic segmentation to consider bbox only annotations
 
 Bug fixes
 ^^^^^^^^^

src/otx/__init__.py

@@ -3,7 +3,7 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-__version__ = "2.2.0rc3"
+__version__ = "2.2.0rc4"
 
 import os
 from pathlib import Path

src/otx/core/data/dataset/segmentation.py

@@ -11,7 +11,7 @@
 import cv2
 import numpy as np
 import torch
-from datumaro.components.annotation import Ellipse, Image, Mask, Polygon
+from datumaro.components.annotation import Bbox, Ellipse, Image, Mask, Polygon, RotatedBbox
 from torchvision import tv_tensors
 
 from otx.core.data.dataset.base import Transforms
@@ -99,11 +99,11 @@ def _extract_class_mask(item: DatasetItem, img_shape: tuple[int, int], ignore_in
         raise ValueError(msg, ignore_index)
 
     # fill mask with background label if we have Polygon/Ellipse annotations
-    fill_value = 0 if isinstance(item.annotations[0], (Ellipse, Polygon)) else ignore_index
+    fill_value = 0 if isinstance(item.annotations[0], (Ellipse, Polygon, Bbox, RotatedBbox)) else ignore_index
     class_mask = np.full(shape=img_shape[:2], fill_value=fill_value, dtype=np.uint8)
 
     for mask in sorted(
-        [ann for ann in item.annotations if isinstance(ann, (Mask, Ellipse, Polygon))],
+        [ann for ann in item.annotations if isinstance(ann, (Mask, Ellipse, Polygon, Bbox, RotatedBbox))],
         key=lambda ann: ann.z_order,
     ):
         index = mask.label
@@ -112,7 +112,7 @@ def _extract_class_mask(item: DatasetItem, img_shape: tuple[int, int], ignore_in
             msg = "Mask's label index should not be None."
             raise ValueError(msg)
 
-        if isinstance(mask, (Ellipse, Polygon)):
+        if isinstance(mask, (Ellipse, Polygon, Bbox, RotatedBbox)):
             polygons = np.asarray(mask.as_polygon(), dtype=np.int32).reshape((-1, 1, 2))
             class_index = index + 1  # NOTE: disregard the background index. Objects start from index=1
             this_class_mask = cv2.drawContours(
@@ -193,8 +193,8 @@ def __init__(
     @property
     def has_polygons(self) -> bool:
         """Check if the dataset has polygons in annotations."""
-        ann_types = {str(ann_type).split(".")[-1] for ann_type in self.dm_subset.ann_types()}
-        if ann_types & {"polygon", "ellipse"}:
+        # all polygon-like format should be considered as polygons
+        if {ann_type.name for ann_type in self.dm_subset.ann_types()} & {"polygon", "ellipse", "bbox", "rotated_bbox"}:
             return True
         return False
 

tests/unit/core/data/conftest.py

@@ -10,7 +10,7 @@
 import cv2
 import numpy as np
 import pytest
-from datumaro.components.annotation import Bbox, Label, LabelCategories, Mask, Polygon
+from datumaro.components.annotation import AnnotationType, Bbox, Label, LabelCategories, Mask, Polygon
 from datumaro.components.dataset import Dataset as DmDataset
 from datumaro.components.dataset_base import DatasetItem
 from datumaro.components.media import Image
@@ -89,7 +89,7 @@ def fxt_dm_item(request, tmpdir) -> DatasetItem:
         media=media,
         annotations=[
             Label(label=0),
-            Bbox(x=0, y=0, w=1, h=1, label=0),
+            Bbox(x=200, y=200, w=1, h=1, label=0),
             Mask(label=0, image=np.eye(10, dtype=np.uint8)),
             Polygon(points=[399.0, 570.0, 397.0, 572.0, 397.0, 573.0, 394.0, 576.0], label=0),
         ],
@@ -133,6 +133,12 @@ def fxt_mock_dm_subset(mocker: MockerFixture, fxt_dm_item: DatasetItem) -> Magic
     mock_dm_subset.__getitem__.return_value = fxt_dm_item
     mock_dm_subset.__len__.return_value = 1
     mock_dm_subset.categories().__getitem__.return_value = LabelCategories.from_iterable(_LABEL_NAMES)
+    mock_dm_subset.ann_types.return_value = [
+        AnnotationType.label,
+        AnnotationType.bbox,
+        AnnotationType.mask,
+        AnnotationType.polygon,
+    ]
     return mock_dm_subset
 
 
@@ -142,6 +148,7 @@ def fxt_mock_det_dm_subset(mocker: MockerFixture, fxt_dm_item_bbox_only: Dataset
     mock_dm_subset.__getitem__.return_value = fxt_dm_item_bbox_only
     mock_dm_subset.__len__.return_value = 1
     mock_dm_subset.categories().__getitem__.return_value = LabelCategories.from_iterable(_LABEL_NAMES)
+    mock_dm_subset.ann_types.return_value = [AnnotationType.bbox]
     return mock_dm_subset
 
 

tests/unit/core/data/dataset/test_segmentation.py

@@ -0,0 +1,36 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+"""Unit tests of classification datasets."""
+
+from otx.core.data.dataset.segmentation import OTXSegmentationDataset
+from otx.core.data.entity.segmentation import SegDataEntity
+
+
+class TestOTXSegmentationDataset:
+    def test_get_item(
+        self,
+        fxt_mock_dm_subset,
+    ) -> None:
+        dataset = OTXSegmentationDataset(
+            dm_subset=fxt_mock_dm_subset,
+            transforms=[lambda x: x],
+            mem_cache_img_max_size=None,
+            max_refetch=3,
+        )
+        assert isinstance(dataset[0], SegDataEntity)
+        assert "background" in [label_name.lower() for label_name in dataset.label_info.label_names]
+
+    def test_get_item_from_bbox_dataset(
+        self,
+        fxt_mock_det_dm_subset,
+    ) -> None:
+        dataset = OTXSegmentationDataset(
+            dm_subset=fxt_mock_det_dm_subset,
+            transforms=[lambda x: x],
+            mem_cache_img_max_size=None,
+            max_refetch=3,
+        )
+        assert isinstance(dataset[0], SegDataEntity)
+        # OTXSegmentationDataset should add background when getting a dataset which includes only bbox annotations
+        assert "background" in [label_name.lower() for label_name in dataset.label_info.label_names]

tox.ini

@@ -55,7 +55,7 @@ commands =
         {posargs}
 
 
-[testenv:integration-test-{all, action, classification, multi_cls_classification, multi_label_classification, hlabel_classification, detection, rotated_detection, keypoint_detection, instance_segmentation, semantic_segmentation, visual_prompting_all, visual_prompting, zero_shot_visual_prompting, anomaly_classification, anomaly_detection, anomaly_segmentation}]
+[testenv:integration-test-{all, action, classification, multi_cls_classification, multi_label_classification, hlabel_classification, detection, rotated_detection, keypoint_detection, instance_segmentation, semantic_segmentation, visual_prompting_all, visual_prompting, zero_shot_visual_prompting, anomaly, anomaly_classification, anomaly_detection, anomaly_segmentation}]
 setenv =
     CUBLAS_WORKSPACE_CONFIG=:4096:8
 deps =
@@ -75,7 +75,7 @@ commands =
 
 [testenv:perf-benchmark]
 deps =
-    .[base,dev]
+    .[base,dev,ci_benchmark]
 commands =
     pytest -ra --showlocals --csv={toxworkdir}/{envname}-test.csv {posargs:tests/perf}