Fix label info dispatching (#4443)

* fix loading checkpoints for the fine tuning

* add test to check label dispatch

Author: kprokofi

src/otx/core/model/base.py

@@ -13,7 +13,7 @@
 import warnings
 from abc import abstractmethod
 from dataclasses import asdict
-from typing import TYPE_CHECKING, Any, Callable, Literal, Sequence
+from typing import TYPE_CHECKING, Any, Callable, Literal
 
 import numpy as np
 import openvino
@@ -410,8 +410,7 @@ def load_state_dict_incrementally(self, ckpt: dict[str, Any], *args, **kwargs) -
             msg = "Checkpoint should have `label_info`."
             raise ValueError(msg, ckpt_label_info)
 
-        if isinstance(ckpt_label_info, dict):
-            ckpt_label_info = LabelInfo(**ckpt_label_info)
+        ckpt_label_info = self._dispatch_label_info(ckpt_label_info)
 
         if not hasattr(ckpt_label_info, "label_ids"):
             msg = "Loading checkpoint from OTX < 2.2.1, label_ids are assigned automatically"
@@ -840,7 +839,7 @@ def _dispatch_label_info(label_info: LabelInfoTypes) -> LabelInfo:
             return LabelInfo(**label_info)
         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):
+        if isinstance(label_info, (list, tuple)) and all(isinstance(name, str) for name in label_info):
             return LabelInfo(
                 label_names=label_info,
                 label_groups=[label_info],

src/otx/core/model/segmentation.py

@@ -29,7 +29,7 @@
 from otx.core.model.base import DefaultOptimizerCallable, DefaultSchedulerCallable, OTXModel, OVModel
 from otx.core.schedulers import LRSchedulerListCallable
 from otx.core.types.export import TaskLevelExportParameters
-from otx.core.types.label import LabelInfo, LabelInfoTypes, SegLabelInfo
+from otx.core.types.label import LabelInfoTypes, SegLabelInfo
 from otx.core.utils.tile_merge import SegmentationTileMerge
 
 if TYPE_CHECKING:
@@ -192,7 +192,7 @@ def _convert_pred_entity_to_compute_metric(
         ]
 
     @staticmethod
-    def _dispatch_label_info(label_info: LabelInfoTypes) -> LabelInfo:
+    def _dispatch_label_info(label_info: LabelInfoTypes) -> SegLabelInfo:
         if isinstance(label_info, dict):
             if "label_ids" not in label_info:
                 # NOTE: This is for backward compatibility

src/otx/core/types/label.py

@@ -420,7 +420,7 @@ class SegLabelInfo(LabelInfo):
     ignore_index: int = 255
 
     @classmethod
-    def from_num_classes(cls, num_classes: int) -> LabelInfo:
+    def from_num_classes(cls, num_classes: int) -> SegLabelInfo:
         """Create this object from the number of classes.
 
         Args:
@@ -437,7 +437,7 @@ def from_num_classes(cls, num_classes: int) -> LabelInfo:
             label_names = ["background", "label_0"]
             return SegLabelInfo(label_names=label_names, label_groups=[label_names], label_ids=["0", "1"])
 
-        return super().from_num_classes(num_classes)
+        return super().from_num_classes(num_classes)  # type: ignore[return-value]
 
 
 @dataclass

tests/unit/core/model/test_base.py

@@ -7,14 +7,19 @@
 import openvino as ov
 import pytest
 import torch
+from datumaro import LabelCategories
+from datumaro.components.annotation import GroupType
 from lightning import Trainer
 from lightning.pytorch.utilities.types import LRSchedulerConfig
 from model_api.models.result import ClassificationResult
 from pytest_mock import MockerFixture
 
 from otx.core.data.entity.base import OTXBatchDataEntity
 from otx.core.model.base import OTXModel, OVModel
+from otx.core.model.classification import OTXHlabelClsModel, OTXMulticlassClsModel
+from otx.core.model.segmentation import OTXSegmentationModel
 from otx.core.schedulers.warmup_schedulers import LinearWarmupScheduler
+from otx.core.types.label import HLabelInfo, LabelInfo, SegLabelInfo
 from tests.unit.core.utils.test_utils import get_dummy_ov_cls_model
 
 
@@ -83,6 +88,108 @@ def test_smart_weight_loading(self, mocker) -> None:
             prev_state_dict["model.head.bias"],
         )
 
+    def test_label_info_dispatch(self, mocker):
+        with mocker.patch.object(OTXModel, "_create_model", return_value=MockNNModule(3)):
+            with pytest.raises(TypeError, match="invalid_label_info"):
+                OTXModel(label_info="invalid_label_info")
+
+            # Test with LabelInfo
+            label_info = OTXModel(
+                label_info=LabelInfo(
+                    ["label_1", "label_2"],
+                    label_ids=["1", "2"],
+                    label_groups=[["label_1", "label_2"]],
+                ),
+            )
+            assert isinstance(label_info.label_info, LabelInfo)
+
+            # Test with SegLabelInfo
+            seg_label_info = OTXModel(label_info=SegLabelInfo.from_num_classes(3))
+            assert isinstance(seg_label_info.label_info, SegLabelInfo)
+
+        with mocker.patch.object(OTXMulticlassClsModel, "_create_model", return_value=MockNNModule(3)):
+            # Test simple Classfication model loading checkpoint
+            cls_model = OTXMulticlassClsModel(
+                label_info=LabelInfo(
+                    ["label_1", "label_2"],
+                    label_ids=["1", "2"],
+                    label_groups=[["label_1", "label_2"]],
+                ),
+                input_size=(224, 224),
+            )
+            label_info_dict = {
+                "label_ids": ["1", "2"],
+                "label_names": ["label_1", "label_2"],
+                "label_groups": [["label_1", "label_2"]],
+            }
+            cls_model.load_state_dict_incrementally(
+                {"state_dict": cls_model.state_dict(), "hyper_parameters": {"label_info": label_info_dict}},
+            )
+            assert isinstance(cls_model.label_info, LabelInfo)
+            # test if ignore_index is not set
+            label_info_dict["ignore_index"] = 255
+            with pytest.raises(TypeError, match=r"unexpected keyword argument.*ignore_index"):
+                cls_model.load_state_dict_incrementally(
+                    {"state_dict": cls_model.state_dict(), "hyper_parameters": {"label_info": label_info_dict}},
+                )
+
+        with mocker.patch.object(OTXSegmentationModel, "_create_model", return_value=MockNNModule(3)):
+            # test segmentation model loading checkpoint with SegLabelInfo
+            segmentation_model = OTXSegmentationModel(
+                label_info=SegLabelInfo.from_num_classes(3),
+                input_size=(224, 224),
+                model_name="segmentation_model",
+            )
+            segmentation_model.load_state_dict_incrementally(
+                {"state_dict": segmentation_model.state_dict(), "hyper_parameters": {"label_info": label_info_dict}},
+            )
+            assert isinstance(segmentation_model.label_info, SegLabelInfo)
+            assert hasattr(segmentation_model.label_info, "ignore_index")
+            assert segmentation_model.label_info.ignore_index == 255
+
+        # test hlabel classification model loading checkpoint with HLabelInfo
+        labels = [
+            LabelCategories.Category(name="car", parent="vehicle"),
+            LabelCategories.Category(name="truck", parent="vehicle"),
+            LabelCategories.Category(name="plush toy", parent="plush toy"),
+            LabelCategories.Category(name="No class"),
+        ]
+        label_groups = [
+            LabelCategories.LabelGroup(
+                name="Detection labels___vehicle",
+                labels=["car", "truck"],
+                group_type=GroupType.EXCLUSIVE,
+            ),
+            LabelCategories.LabelGroup(
+                name="Detection labels___plush toy",
+                labels=["plush toy"],
+                group_type=GroupType.EXCLUSIVE,
+            ),
+            LabelCategories.LabelGroup(name="No class", labels=["No class"], group_type=GroupType.RESTRICTED),
+        ]
+        dm_label_categories = LabelCategories(items=labels, label_groups=label_groups)
+        hlabel_info = HLabelInfo.from_dm_label_groups(dm_label_categories)
+        hlabel_dict_label_info = hlabel_info.as_dict(normalize_label_names=True)
+
+        with mocker.patch.object(OTXHlabelClsModel, "_create_model", return_value=MockNNModule(3)):
+            hlabel_model = OTXHlabelClsModel(hlabel_dict_label_info, input_size=(224, 224))
+            hlabel_model.load_state_dict_incrementally(
+                {"state_dict": hlabel_model.state_dict(), "hyper_parameters": {"label_info": hlabel_dict_label_info}},
+            )
+
+            with pytest.raises(TypeError, match=r"unexpected keyword argument.*num_multiclass_heads"):
+                segmentation_model.load_state_dict_incrementally(
+                    {
+                        "state_dict": segmentation_model.state_dict(),
+                        "hyper_parameters": {"label_info": hlabel_dict_label_info},
+                    },
+                )
+
+            with pytest.raises(TypeError, match=r"unexpected keyword argument.*num_multiclass_heads"):
+                cls_model.load_state_dict_incrementally(
+                    {"state_dict": cls_model.state_dict(), "hyper_parameters": {"label_info": hlabel_dict_label_info}},
+                )
+
     def test_lr_scheduler_step(self, mocker: MockerFixture) -> None:
         mock_linear_warmup_scheduler = mocker.create_autospec(spec=LinearWarmupScheduler)
         mock_main_scheduler = mocker.create_autospec(spec=torch.optim.lr_scheduler.LRScheduler)