🚀 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, trainer  # 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]
@@ -182,16 +181,37 @@ 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 name, metric in pixel_metric.items(copy_state=False):
+            if isinstance(metric, SPRO):
+                metric.update(torch.squeeze(output["anomaly_maps"]), output["masks"])
+            else:
+                logger.warning(
+                    f"Metric {name} may not be suitable for a dataset with the region separated "
+                    "in multiple ground-truth masks.",
+                )
+                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/cli/cli.py

@@ -143,7 +143,12 @@ def add_arguments_to_parser(self, parser: LightningArgumentParser) -> None:
         parser.add_argument("--visualization.show", type=bool, default=False)
         parser.add_argument("--task", type=TaskType, default=TaskType.SEGMENTATION)
         parser.add_argument("--metrics.image", type=list[str] | str | None, default=["F1Score", "AUROC"])
-        parser.add_argument("--metrics.pixel", type=list[str] | str | None, default=None, required=False)
+        parser.add_argument(
+            "--metrics.pixel",
+            type=list[str] | str | dict[str, dict[str, Any]] | None,
+            default=None,
+            required=False,
+        )
         parser.add_argument("--metrics.threshold", type=BaseThreshold, default="F1AdaptiveThreshold")
         parser.add_argument("--logging.log_graph", type=bool, help="Log the model to the logger", default=False)
         if hasattr(parser, "subcommand") and parser.subcommand != "predict":  # Predict also accepts str and Path inputs

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"]