Update label info (#3925)

add label info to init

Signed-off-by: Ashwin Vaidya <[email protected]>

Author: ashwinvaidya17

src/otx/algo/anomaly/padim.py

@@ -14,13 +14,15 @@
 from anomalib.models.image import Padim as AnomalibPadim
 
 from otx.core.model.anomaly import OTXAnomaly
+from otx.core.types.label import AnomalyLabelInfo
 from otx.core.types.task import OTXTaskType
 
 if TYPE_CHECKING:
     from lightning.pytorch.utilities.types import STEP_OUTPUT
     from torch.optim.optimizer import Optimizer
 
     from otx.core.model.anomaly import AnomalyModelInputs, AnomalyModelOutputs
+    from otx.core.types.label import LabelInfoTypes
 
 
 class Padim(OTXAnomaly, AnomalibPadim):
@@ -40,6 +42,7 @@ class Padim(OTXAnomaly, AnomalibPadim):
 
     def __init__(
         self,
+        label_info: LabelInfoTypes = AnomalyLabelInfo(),
         backbone: str = "resnet18",
         layers: list[str] = ["layer1", "layer2", "layer3"],  # noqa: B006
         pre_trained: bool = True,
@@ -51,7 +54,7 @@ def __init__(
         ] = OTXTaskType.ANOMALY_CLASSIFICATION,
         input_size: tuple[int, int] = (256, 256),
     ) -> None:
-        OTXAnomaly.__init__(self, input_size)
+        OTXAnomaly.__init__(self, label_info=label_info, input_size=input_size)
         AnomalibPadim.__init__(
             self,
             backbone=backbone,

src/otx/algo/anomaly/stfpm.py

@@ -14,13 +14,15 @@
 from anomalib.models.image.stfpm import Stfpm as AnomalibStfpm
 
 from otx.core.model.anomaly import OTXAnomaly
+from otx.core.types.label import AnomalyLabelInfo
 from otx.core.types.task import OTXTaskType
 
 if TYPE_CHECKING:
     from lightning.pytorch.utilities.types import STEP_OUTPUT
     from torch.optim.optimizer import Optimizer
 
     from otx.core.model.anomaly import AnomalyModelInputs, AnomalyModelOutputs
+    from otx.core.types.label import LabelInfoTypes
 
 
 class Stfpm(OTXAnomaly, AnomalibStfpm):
@@ -38,6 +40,7 @@ class Stfpm(OTXAnomaly, AnomalibStfpm):
 
     def __init__(
         self,
+        label_info: LabelInfoTypes = AnomalyLabelInfo(),
         layers: Sequence[str] = ["layer1", "layer2", "layer3"],
         backbone: str = "resnet18",
         task: Literal[
@@ -48,7 +51,7 @@ def __init__(
         input_size: tuple[int, int] = (256, 256),
         **kwargs,
     ) -> None:
-        OTXAnomaly.__init__(self, input_size=input_size)
+        OTXAnomaly.__init__(self, label_info=label_info, input_size=input_size)
         AnomalibStfpm.__init__(
             self,
             backbone=backbone,

src/otx/core/model/anomaly.py

@@ -39,7 +39,8 @@
     from lightning.pytorch.callbacks.callback import Callback
     from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable
     from torchmetrics import Metric
-from otx.core.types.label import AnomalyLabelInfo
+
+    from otx.core.types.label import LabelInfoTypes
 
 AnomalyModelInputs: TypeAlias = (
     AnomalyClassificationDataBatch | AnomalySegmentationDataBatch | AnomalyDetectionDataBatch
@@ -57,8 +58,8 @@ class OTXAnomaly(OTXModel):
             Model input size in the order of height and width. Defaults to None.
     """
 
-    def __init__(self, input_size: tuple[int, int]) -> None:
-        super().__init__(label_info=AnomalyLabelInfo(), input_size=input_size)
+    def __init__(self, label_info: LabelInfoTypes, input_size: tuple[int, int]) -> None:
+        super().__init__(label_info=label_info, input_size=input_size)
         self.optimizer: list[OptimizerCallable] | OptimizerCallable = None
         self.scheduler: list[LRSchedulerCallable] | LRSchedulerCallable = None
         self.trainer: Trainer