Update Label Info handling

src/otx/algo/classification/mobilenet_v3.py

@@ -15,7 +15,7 @@
 from otx.algo.classification.backbones import OTXMobileNetV3
 from otx.algo.classification.classifier import HLabelClassifier, ImageClassifier, SemiSLClassifier
 from otx.algo.classification.heads import (
-    HierarchicalCBAMClsHead,
+    HierarchicalLinearClsHead,
     LinearClsHead,
     MultiLabelNonLinearClsHead,
     SemiSLLinearClsHead,
@@ -313,14 +313,12 @@ def _build_model(self, head_config: dict) -> nn.Module:
 
         copied_head_config = copy(head_config)
         copied_head_config["step_size"] = (ceil(self.input_size[0] / 32), ceil(self.input_size[1] / 32))
+        in_channels = 960 if self.mode == "large" else 576
 
         return HLabelClassifier(
             backbone=OTXMobileNetV3(mode=self.mode, input_size=self.input_size),
-            neck=nn.Identity(),
-            head=HierarchicalCBAMClsHead(
-                in_channels=960,
-                **copied_head_config,
-            ),
+            neck=GlobalAveragePooling(dim=2),
+            head=HierarchicalLinearClsHead(**copied_head_config, in_channels=in_channels),
             multiclass_loss=nn.CrossEntropyLoss(),
             multilabel_loss=AsymmetricAngularLossWithIgnore(gamma_pos=0.0, gamma_neg=1.0, reduction="sum"),
         )

src/otx/algo/classification/timm_model.py

@@ -15,12 +15,12 @@
 from otx.algo.classification.backbones.timm import TimmBackbone, TimmModelType
 from otx.algo.classification.classifier import HLabelClassifier, ImageClassifier, SemiSLClassifier
 from otx.algo.classification.heads import (
-    HierarchicalCBAMClsHead,
     LinearClsHead,
     MultiLabelLinearClsHead,
     SemiSLLinearClsHead,
 )
 from otx.algo.classification.losses.asymmetric_angular_loss_with_ignore import AsymmetricAngularLossWithIgnore
+from otx.algo.classification.mobilenet_v3 import HierarchicalLinearClsHead
 from otx.algo.classification.necks.gap import GlobalAveragePooling
 from otx.algo.classification.utils import get_classification_layers
 from otx.algo.utils.support_otx_v1 import OTXv1Helper
@@ -272,11 +272,8 @@ def _build_model(self, head_config: dict) -> nn.Module:
         copied_head_config["step_size"] = (ceil(self.input_size[0] / 32), ceil(self.input_size[1] / 32))
         return HLabelClassifier(
             backbone=backbone,
-            neck=nn.Identity(),
-            head=HierarchicalCBAMClsHead(
-                in_channels=backbone.num_features,
-                **copied_head_config,
-            ),
+            neck=GlobalAveragePooling(dim=2),
+            head=HierarchicalLinearClsHead(**copied_head_config, in_channels=backbone.num_features),
             multiclass_loss=nn.CrossEntropyLoss(),
             multilabel_loss=AsymmetricAngularLossWithIgnore(gamma_pos=0.0, gamma_neg=1.0, reduction="sum"),
         )

src/otx/algo/classification/torchvision_model.py

@@ -14,12 +14,12 @@
 from otx.algo.classification.backbones.torchvision import TorchvisionBackbone, TVModelType
 from otx.algo.classification.classifier import HLabelClassifier, ImageClassifier, SemiSLClassifier
 from otx.algo.classification.heads import (
-    HierarchicalCBAMClsHead,
     LinearClsHead,
     MultiLabelLinearClsHead,
     SemiSLLinearClsHead,
 )
 from otx.algo.classification.losses import AsymmetricAngularLossWithIgnore
+from otx.algo.classification.mobilenet_v3 import HierarchicalLinearClsHead
 from otx.algo.classification.necks.gap import GlobalAveragePooling
 from otx.algo.classification.utils import get_classification_layers
 from otx.core.data.entity.classification import (
@@ -315,11 +315,8 @@ def _build_model(self, head_config: dict) -> nn.Module:
         backbone = TorchvisionBackbone(backbone=self.backbone, pretrained=self.pretrained)
         return HLabelClassifier(
             backbone=backbone,
-            neck=nn.Identity(),
-            head=HierarchicalCBAMClsHead(
-                in_channels=backbone.in_features,
-                **head_config,
-            ),
+            neck=GlobalAveragePooling(dim=2),
+            head=HierarchicalLinearClsHead(**head_config, in_channels=backbone.in_features),
             multiclass_loss=nn.CrossEntropyLoss(),
             multilabel_loss=AsymmetricAngularLossWithIgnore(gamma_pos=0.0, gamma_neg=1.0, reduction="sum"),
         )

src/otx/algo/classification/vit.py

@@ -19,12 +19,12 @@
 from otx.algo.classification.backbones.vision_transformer import VIT_ARCH_TYPE, VisionTransformer
 from otx.algo.classification.classifier import HLabelClassifier, ImageClassifier, SemiSLClassifier
 from otx.algo.classification.heads import (
-    HierarchicalCBAMClsHead,
     MultiLabelLinearClsHead,
     SemiSLVisionTransformerClsHead,
     VisionTransformerClsHead,
 )
 from otx.algo.classification.losses import AsymmetricAngularLossWithIgnore
+from otx.algo.classification.mobilenet_v3 import HierarchicalLinearClsHead
 from otx.algo.classification.utils import get_classification_layers
 from otx.algo.explain.explain_algo import ViTReciproCAM, feature_vector_fn
 from otx.algo.utils.support_otx_v1 import OTXv1Helper
@@ -466,11 +466,7 @@ def _build_model(self, head_config: dict) -> nn.Module:
         return HLabelClassifier(
             backbone=vit_backbone,
             neck=None,
-            head=HierarchicalCBAMClsHead(
-                in_channels=vit_backbone.embed_dim,
-                step_size=1,
-                **head_config,
-            ),
+            head=HierarchicalLinearClsHead(**head_config, in_channels=vit_backbone.embed_dim),
             multiclass_loss=nn.CrossEntropyLoss(),
             multilabel_loss=AsymmetricAngularLossWithIgnore(gamma_pos=0.0, gamma_neg=1.0, reduction="sum"),
             init_cfg=init_cfg,

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

@@ -37,6 +37,7 @@ def __init__(
         image_color_channel: ImageColorChannel = ImageColorChannel.BGR,
         stack_images: bool = True,
         to_tv_image: bool = True,
+        data_format: str = "",
     ) -> None:
         super().__init__(
             dm_subset,

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

@@ -57,6 +57,7 @@ def __init__(
         image_color_channel: ImageColorChannel = ImageColorChannel.RGB,
         stack_images: bool = True,
         to_tv_image: bool = True,
+        data_format: str = "",
     ) -> None:
         self.task_type = task_type
         super().__init__(

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

@@ -83,6 +83,7 @@
         image_color_channel: ImageColorChannel = ImageColorChannel.RGB,
         stack_images: bool = True,
         to_tv_image: bool = True,
+        data_format: str = "",
     ) -> None:
         self.dm_subset = dm_subset
         self.transforms = transforms
@@ -92,8 +93,11 @@
         self.image_color_channel = image_color_channel
         self.stack_images = stack_images
         self.to_tv_image = to_tv_image
+        self.data_format = data_format
 
-        if self.dm_subset.categories():
+        if self.dm_subset.categories() and data_format == "arrow":
+            self.label_info = LabelInfo.from_dm_label_groups_arrow(self.dm_subset.categories()[AnnotationType.label])
+        elif self.dm_subset.categories():
             self.label_info = LabelInfo.from_dm_label_groups(self.dm_subset.categories()[AnnotationType.label])
         else:
             self.label_info = NullLabelInfo()

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

@@ -128,7 +128,15 @@
         self.dm_categories = self.dm_subset.categories()[AnnotationType.label]
 
         # Hlabel classification used HLabelInfo to insert the HLabelData.
-        self.label_info = HLabelInfo.from_dm_label_groups(self.dm_categories)
+        if self.data_format == "arrow":
+            # arrow format stores label IDs as names, have to deal with that here
+            self.label_info = HLabelInfo.from_dm_label_groups_arrow(self.dm_categories)
+        else:
+            self.label_info = HLabelInfo.from_dm_label_groups(self.dm_categories)
+
+        self.id_to_name_mapping = dict(zip(self.label_info.label_ids, self.label_info.label_names))
+        self.id_to_name_mapping[""] = ""
+
         if self.label_info.num_multiclass_heads == 0:
             msg = "The number of multiclass heads should be larger than 0."
             raise ValueError(msg)
@@ -149,14 +157,16 @@
         """
 
         def _label_idx_to_name(idx: int) -> str:
-            return self.label_info.label_names[idx]
+            return self.dm_categories[idx].name
 
         def _label_name_to_idx(name: str) -> int:
             indices = [idx for idx, val in enumerate(self.label_info.label_names) if val == name]
             return indices[0]
 
         def _get_label_group_idx(label_name: str) -> int:
             if isinstance(self.label_info, HLabelInfo):
+                if self.data_format == "arrow":
+                    return self.label_info.class_to_group_idx[self.id_to_name_mapping[label_name]][0]
                 return self.label_info.class_to_group_idx[label_name][0]
             msg = f"self.label_info should have HLabelInfo type, got {type(self.label_info)}"
             raise ValueError(msg)
@@ -256,18 +266,18 @@
             class_indices[i] = -1
 
         for ann in label_anns:
-            ann_name = self.dm_categories.items[ann.label].name
-            ann_parent = self.dm_categories.items[ann.label].parent
+            if self.data_format == "arrow":
+                # skips unknown labels for instance, the empty one
+                if self.dm_categories.items[ann.label].name not in self.id_to_name_mapping:
+                    continue
+                ann_name = self.id_to_name_mapping[self.dm_categories.items[ann.label].name]
+            else:
+                ann_name = self.dm_categories.items[ann.label].name
             group_idx, in_group_idx = self.label_info.class_to_group_idx[ann_name]
-            (parent_group_idx, parent_in_group_idx) = (
-                self.label_info.class_to_group_idx[ann_parent] if ann_parent else (None, None)
-            )
 
             if group_idx < num_multiclass_heads:
                 class_indices[group_idx] = in_group_idx
-                if parent_group_idx is not None and parent_in_group_idx is not None:
-                    class_indices[parent_group_idx] = parent_in_group_idx
-            elif not ignored_labels or ann.label not in ignored_labels:
+            elif ann.label not in ignored_labels:
                 class_indices[num_multiclass_heads + in_group_idx] = 1
             else:
                 class_indices[num_multiclass_heads + in_group_idx] = -1

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

@@ -54,9 +54,11 @@ def __init__(
         self.dm_subset = self._get_single_bbox_dataset(dm_subset)
 
         if self.dm_subset.categories():
+            kp_labels = self.dm_subset.categories()[AnnotationType.points][0].labels
             self.label_info = LabelInfo(
-                label_names=self.dm_subset.categories()[AnnotationType.points][0].labels,
+                label_names=kp_labels,
                 label_groups=[],
+                label_ids=[str(i) for i in range(len(kp_labels))],
             )
         else:
             self.label_info = NullLabelInfo()

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

@@ -167,6 +167,7 @@ def __init__(
         stack_images: bool = True,
         to_tv_image: bool = True,
         ignore_index: int = 255,
+        data_format: str = "",
     ) -> None:
         super().__init__(
             dm_subset,
@@ -187,6 +188,7 @@ def __init__(
             label_names=self.label_info.label_names,
             label_groups=self.label_info.label_groups,
             ignore_index=ignore_index,
+            label_ids=self.label_info.label_ids,
         )
         self.ignore_index = ignore_index
 

src/otx/core/data/factory.py

@@ -73,6 +73,7 @@ def create(  # noqa: PLR0911
         dm_subset: DmDataset,
         cfg_subset: SubsetConfig,
         mem_cache_handler: MemCacheHandlerBase,
+        data_format: str,
         mem_cache_img_max_size: tuple[int, int] | None = None,
         image_color_channel: ImageColorChannel = ImageColorChannel.RGB,
         stack_images: bool = True,
@@ -85,6 +86,7 @@ def create(  # noqa: PLR0911
         common_kwargs = {
             "dm_subset": dm_subset,
             "transforms": transforms,
+            "data_format": data_format,
             "mem_cache_handler": mem_cache_handler,
             "mem_cache_img_max_size": mem_cache_img_max_size,
             "image_color_channel": image_color_channel,

src/otx/core/data/module.py

@@ -107,13 +107,6 @@ def __init__(  # noqa: PLR0913
         self.subsets: dict[str, OTXDataset] = {}
         self.save_hyperparameters(ignore=["input_size"])
 
-        # TODO (Jaeguk): This is workaround for a bug in Datumaro.
-        # These lines should be removed after next datumaro release.
-        # https://github.com/openvinotoolkit/datumaro/pull/1223/files
-        from datumaro.plugins.data_formats.video import VIDEO_EXTENSIONS
-
-        VIDEO_EXTENSIONS.append(".mp4")
-
         dataset = DmDataset.import_from(self.data_root, format=self.data_format)
         if self.task != "H_LABEL_CLS":
             dataset = pre_filtering(
@@ -193,6 +186,7 @@ def __init__(  # noqa: PLR0913
                 dm_subset=dm_subset.as_dataset(),
                 cfg_subset=config_mapping[name],
                 mem_cache_handler=mem_cache_handler,
+                data_format=self.data_format,
                 mem_cache_img_max_size=mem_cache_img_max_size,
                 image_color_channel=image_color_channel,
                 stack_images=stack_images,
@@ -237,6 +231,7 @@ def __init__(  # noqa: PLR0913
                         include_polygons=include_polygons,
                         ignore_index=ignore_index,
                         vpm_config=vpm_config,
+                        data_format=self.data_format,
                     )
                     self.subsets[transform_key] = unlabeled_dataset
             else:
@@ -251,6 +246,7 @@ def __init__(  # noqa: PLR0913
                     include_polygons=include_polygons,
                     ignore_index=ignore_index,
                     vpm_config=vpm_config,
+                    data_format=self.data_format,
                 )
                 self.subsets[name] = unlabeled_dataset
 

src/otx/core/data/pre_filtering.py

@@ -84,7 +84,7 @@ def remove_unused_labels(dataset: DmDataset, data_format: str, ignore_index: int
         used_labels = [0, *used_labels]
     if data_format == "common_semantic_segmentation_with_subset_dirs" and len(original_categories) < len(used_labels):
         msg = (
-            "There are labeles mismatch in dataset categories and actuall categories comes from semantic masks."
+            "There are labels mismatch in dataset categories and actual categories comes from semantic masks."
             "Please, check `dataset_meta.json` file."
         )
         raise ValueError(msg)

src/otx/core/model/base.py

@@ -809,7 +809,11 @@ def _dispatch_label_info(label_info: LabelInfoTypes) -> LabelInfo:
         if isinstance(label_info, int):
             return LabelInfo.from_num_classes(num_classes=label_info)
         if isinstance(label_info, Sequence) and all(isinstance(name, str) for name in label_info):
-            return LabelInfo(label_names=label_info, label_groups=[label_info])
+            return LabelInfo(
+                label_names=label_info,
+                label_groups=[label_info],
+                label_ids=[str(i) for i in range(len(label_info))],
+            )
         if isinstance(label_info, LabelInfo):
             return label_info
 
@@ -1113,7 +1117,7 @@ def _create_label_info_from_ov_ir(self) -> LabelInfo:
             )
 
             logger.warning(msg)
-            return LabelInfo(label_names=label_names, label_groups=[label_names])
+            return LabelInfo(label_names=label_names, label_groups=[label_names], label_ids=[])
 
         msg = "Cannot construct LabelInfo from OpenVINO IR. Please check this model is trained by OTX."
         raise ValueError(msg)

src/otx/core/model/segmentation.py

@@ -245,7 +245,11 @@
         if isinstance(label_info, int):
             return SegLabelInfo.from_num_classes(num_classes=label_info)
         if isinstance(label_info, Sequence) and all(isinstance(name, str) for name in label_info):
-            return SegLabelInfo(label_names=label_info, label_groups=[label_info])
+            return SegLabelInfo(
+                label_names=label_info,
+                label_groups=[label_info],
+                label_ids=[str(i) for i in range(len(label_info))],
+            )
         if isinstance(label_info, SegLabelInfo):
             return label_info
 

src/otx/core/types/export.py

@@ -104,13 +104,15 @@ def to_metadata(self) -> dict[tuple[str, str], str]:
             all_label_ids = "None "
             for lbl in self.label_info.label_names:
                 all_labels += lbl.replace(" ", "_") + " "
-                all_label_ids += lbl.replace(" ", "_") + " "
+            for lbl_id in self.label_info.label_ids:
+                all_label_ids += lbl_id + " "
         else:
             all_labels = ""
             all_label_ids = ""
             for lbl in self.label_info.label_names:
                 all_labels += lbl.replace(" ", "_") + " "
-                all_label_ids += lbl.replace(" ", "_") + " "
+            for lbl_id in self.label_info.label_ids:
+                all_label_ids += lbl_id + " "
 
         metadata = {
             # Common