🚀 Add support for MVTec LOCO dataset and sPRO metric

CHANGELOG.md

@@ -8,6 +8,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Add support for MVTec LOCO AD dataset and sPRO metric by @willyfh in https://github.com/openvinotoolkit/anomalib/pull/1686
+
 ### Changed
 
 - Changed default inference device to AUTO in https://github.com/openvinotoolkit/anomalib/pull/1534

docs/source/markdown/guides/reference/data/image/index.md

@@ -30,6 +30,13 @@ Learn more about Kolektor dataset.
 Learn more about MVTec 2D dataset
 :::
 
+:::{grid-item-card} MVTec LOCO
+:link: ./mvtec_loco
+:link-type: doc
+
+Learn more about MVTec LOCO dataset
+:::
+
 :::{grid-item-card} Visa
 :link: ./visa
 :link-type: doc
@@ -47,5 +54,6 @@ Learn more about Visa dataset.
 ./folder
 ./kolektor
 ./mvtec
+./mvtec_loco
 ./visa
 ```

docs/source/markdown/guides/reference/data/image/mvtec_loco.md

@@ -0,0 +1,7 @@
+# MVTec LOCO Data
+
+```{eval-rst}
+.. automodule:: anomalib.data.image.mvtec_loco
+   :members:
+   :show-inheritance:
+```

src/anomalib/callbacks/metrics.py

@@ -13,7 +13,7 @@
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 
 from anomalib import TaskType
-from anomalib.metrics import AnomalibMetricCollection, create_metric_collection
+from anomalib.metrics import SPRO, AnomalibMetricCollection, create_metric_collection
 from anomalib.models import AnomalyModule
 
 logger = logging.getLogger(__name__)
@@ -67,8 +67,7 @@ def setup(
             pl_module (AnomalyModule): Anomalib Model that inherits pl LightningModule.
             stage (str | None, optional): fit, validate, test or predict. Defaults to None.
         """
-        del trainer, stage  # These variables are not used.
-
+        del stage  # this variable is not used.
         image_metric_names = [] if self.image_metric_names is None else self.image_metric_names
         if isinstance(image_metric_names, str):
             image_metric_names = [image_metric_names]
@@ -98,6 +97,8 @@ def setup(
             else:
                 pl_module.pixel_metrics = create_metric_collection(pixel_metric_names, "pixel_")
             self._set_threshold(pl_module)
+            if hasattr(trainer.datamodule, "saturation_config"):
+                self._set_saturation_config(pl_module, trainer.datamodule.saturation_config)
 
     def on_validation_epoch_start(
         self,
@@ -172,6 +173,9 @@ def _set_threshold(self, pl_module: AnomalyModule) -> None:
         pl_module.image_metrics.set_threshold(pl_module.image_threshold.value.item())
         pl_module.pixel_metrics.set_threshold(pl_module.pixel_threshold.value.item())
 
+    def _set_saturation_config(self, pl_module: AnomalyModule, saturation_config: dict[int, Any]) -> None:
+        pl_module.pixel_metrics.set_saturation_config(saturation_config)
+
     def _update_metrics(
         self,
         image_metric: AnomalibMetricCollection,
@@ -182,16 +186,33 @@ def _update_metrics(
         image_metric.update(output["pred_scores"], output["label"].int())
         if "mask" in output and "anomaly_maps" in output:
             pixel_metric.to(self.device)
-            pixel_metric.update(torch.squeeze(output["anomaly_maps"]), torch.squeeze(output["mask"].int()))
+            if "masks" in output:
+                self._update_pixel_metrics(pixel_metric, output)
+            else:
+                pixel_metric.update(torch.squeeze(output["anomaly_maps"]), torch.squeeze(output["mask"].int()))
 
     def _outputs_to_device(self, output: STEP_OUTPUT) -> STEP_OUTPUT | dict[str, Any]:
         if isinstance(output, dict):
             for key, value in output.items():
                 output[key] = self._outputs_to_device(value)
         elif isinstance(output, torch.Tensor):
             output = output.to(self.device)
+        elif isinstance(output, list):
+            for i, value in enumerate(output):
+                output[i] = self._outputs_to_device(value)
         return output
 
+    def _update_pixel_metrics(self, pixel_metric: AnomalibMetricCollection, output: STEP_OUTPUT) -> None:
+        """Handle metric updates when the SPRO metric is used alongside other pixel-level metrics."""
+        update = False
+        for metric in pixel_metric.values(copy_state=False):
+            if isinstance(metric, SPRO):
+                metric.update(torch.squeeze(output["anomaly_maps"]), output["masks"])
+            else:
+                metric.update(torch.squeeze(output["anomaly_maps"]), torch.squeeze(output["mask"].int()))
+            update = True
+        pixel_metric.set_update_called(update)
+
     @staticmethod
     def _log_metrics(pl_module: AnomalyModule) -> None:
         """Log computed performance metrics."""

src/anomalib/data/__init__.py

@@ -15,7 +15,7 @@
 
 from .base import AnomalibDataModule, AnomalibDataset
 from .depth import DepthDataFormat, Folder3D, MVTec3D
-from .image import BTech, Folder, ImageDataFormat, Kolektor, MVTec, Visa
+from .image import BTech, Folder, ImageDataFormat, Kolektor, MVTec, MVTecLoco, Visa
 from .predict import PredictDataset
 from .video import Avenue, ShanghaiTech, UCSDped, VideoDataFormat
 
@@ -62,6 +62,7 @@ def get_datamodule(config: DictConfig | ListConfig) -> AnomalibDataModule:
     "Kolektor",
     "MVTec",
     "MVTec3D",
+    "MVTecLoco",
     "Avenue",
     "UCSDped",
     "ShanghaiTech",

src/anomalib/data/base/datamodule.py

@@ -26,7 +26,9 @@
 def collate_fn(batch: list) -> dict[str, Any]:
     """Collate bounding boxes as lists.
 
-    Bounding boxes are collated as a list of tensors, while the default collate function is used for all other entries.
+    Bounding boxes and `masks` (not `mask`) are collated as a list of tensors. If `masks` is exist,
+    the `mask_path` is also collated as a list since each element in the batch could be unequal.
+    For all other entries, the default collate function is used.
 
     Args:
         batch (List): list of items in the batch where len(batch) is equal to the batch size.
@@ -40,6 +42,10 @@ def collate_fn(batch: list) -> dict[str, Any]:
         if "boxes" in elem:
             # collate boxes as list
             out_dict["boxes"] = [item.pop("boxes") for item in batch]
+        if "masks" in elem:
+            # collate masks and mask_path as list
+            out_dict["masks"] = [item.pop("masks") for item in batch]
+            out_dict["mask_path"] = [item.pop("mask_path") for item in batch]
         # collate other data normally
         out_dict.update({key: default_collate([item[key] for item in batch]) for key in elem})
         return out_dict
@@ -164,6 +170,9 @@ def _create_val_split(self) -> None:
             # converted from random training sample
             self.train_data, normal_val_data = random_split(self.train_data, self.val_split_ratio, seed=self.seed)
             self.val_data = SyntheticAnomalyDataset.from_dataset(normal_val_data)
+        elif self.val_split_mode == ValSplitMode.FROM_DIR:
+            # the val_data is prepared in subclass
+            pass
         elif self.val_split_mode != ValSplitMode.NONE:
             msg = f"Unknown validation split mode: {self.val_split_mode}"
             raise ValueError(msg)

src/anomalib/data/image/__init__.py

@@ -13,6 +13,7 @@
 from .folder import Folder
 from .kolektor import Kolektor
 from .mvtec import MVTec
+from .mvtec_loco import MVTecLoco
 from .visa import Visa
 
 
@@ -21,11 +22,12 @@ class ImageDataFormat(str, Enum):
 
     MVTEC = "mvtec"
     MVTEC_3D = "mvtec_3d"
+    MVTEC_LOCO = "mvtec_loco"
     BTECH = "btech"
     KOLEKTOR = "kolektor"
     FOLDER = "folder"
     FOLDER_3D = "folder_3d"
     VISA = "visa"
 
 
-__all__ = ["BTech", "Folder", "Kolektor", "MVTec", "Visa"]
+__all__ = ["BTech", "Folder", "Kolektor", "MVTec", "MVTecLoco", "Visa"]