🚀 Add MVTec LOCO dataset and sPRO metric

CHANGELOG.md

@@ -8,6 +8,7 @@ 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
 - 🚀 Update OpenVINO and ONNX export to support fixed input shape by @adrianboguszewski in https://github.com/openvinotoolkit/anomalib/pull/2006
 - Add data_path argument to predict entrypoint and add properties for retrieving model path by @djdameln in https://github.com/openvinotoolkit/anomalib/pull/2018
 

configs/README.md

@@ -15,6 +15,7 @@ configs/
 │   ├── kolektor.yaml
 │   ├── mvtec_3d.yaml
 │   ├── mvtec.yaml
+│   ├── mvtec_loco.yaml
 │   ├── shanghaitec.yaml
 │   ├── ucsd_ped.yaml
 │   └── visa.yaml

configs/data/mvtec_loco.yaml

@@ -0,0 +1,13 @@
+class_path: anomalib.data.MVTecLoco
+init_args:
+  root: ./datasets/MVTec_LOCO
+  category: breakfast_box
+  train_batch_size: 32
+  eval_batch_size: 32
+  num_workers: 8
+  task: SEGMENTATION
+  test_split_mode: FROM_DIR
+  test_split_ratio: 0.2
+  val_split_mode: FROM_DIR
+  val_split_ratio: 0.5
+  seed: null

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 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]
@@ -85,9 +84,25 @@ def setup(
             )
         else:
             pixel_metric_names = (
-                self.pixel_metric_names if not isinstance(self.pixel_metric_names, str) else [self.pixel_metric_names]
+                self.pixel_metric_names.copy()
+                if not isinstance(self.pixel_metric_names, str)
+                else [self.pixel_metric_names]
             )
 
+        # create a separate metric collection for metrics that operate over the semantic segmentation mask
+        # (segmentation mask with a separate channel for each defect type)
+        semantic_pixel_metric_names: list[str] | dict[str, dict[str, Any]]
+        # currently only SPRO metric is supported as semantic segmentation metric
+        if "SPRO" in pixel_metric_names:
+            if isinstance(pixel_metric_names, list):
+                pixel_metric_names.remove("SPRO")
+                semantic_pixel_metric_names = ["SPRO"]
+            elif isinstance(pixel_metric_names, dict):
+                spro_metric = pixel_metric_names.pop("SPRO")
+                semantic_pixel_metric_names = {"SPRO": spro_metric}
+            else:
+                logger.warning("Unexpected type for pixel_metric_names: %s", type(pixel_metric_names))
+
         if isinstance(pl_module, AnomalyModule):
             pl_module.image_metrics = create_metric_collection(image_metric_names, "image_")
             if hasattr(pl_module, "pixel_metrics"):  # incase metrics are loaded from model checkpoint
@@ -97,6 +112,7 @@ def setup(
                         pl_module.pixel_metrics.add_metrics(new_metrics[name])
             else:
                 pl_module.pixel_metrics = create_metric_collection(pixel_metric_names, "pixel_")
+            pl_module.semantic_pixel_metrics = create_metric_collection(semantic_pixel_metric_names, "pixel_")
             self._set_threshold(pl_module)
 
     def on_validation_epoch_start(
@@ -108,6 +124,7 @@ def on_validation_epoch_start(
 
         pl_module.image_metrics.reset()
         pl_module.pixel_metrics.reset()
+        pl_module.semantic_pixel_metrics.reset()
 
     def on_validation_batch_end(
         self,
@@ -122,7 +139,7 @@ def on_validation_batch_end(
 
         if outputs is not None:
             self._outputs_to_device(outputs)
-            self._update_metrics(pl_module.image_metrics, pl_module.pixel_metrics, outputs)
+            self._update_metrics(pl_module, outputs)
 
     def on_validation_epoch_end(
         self,
@@ -143,6 +160,7 @@ def on_test_epoch_start(
 
         pl_module.image_metrics.reset()
         pl_module.pixel_metrics.reset()
+        pl_module.semantic_pixel_metrics.reset()
 
     def on_test_batch_end(
         self,
@@ -157,7 +175,7 @@ def on_test_batch_end(
 
         if outputs is not None:
             self._outputs_to_device(outputs)
-            self._update_metrics(pl_module.image_metrics, pl_module.pixel_metrics, outputs)
+            self._update_metrics(pl_module, outputs)
 
     def on_test_epoch_end(
         self,
@@ -171,32 +189,38 @@ def on_test_epoch_end(
     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())
+        pl_module.semantic_pixel_metrics.set_threshold(pl_module.pixel_threshold.value.item())
 
     def _update_metrics(
         self,
-        image_metric: AnomalibMetricCollection,
-        pixel_metric: AnomalibMetricCollection,
+        pl_module: AnomalyModule,
         output: STEP_OUTPUT,
     ) -> None:
-        image_metric.to(self.device)
-        image_metric.update(output["pred_scores"], output["label"].int())
+        pl_module.image_metrics.to(self.device)
+        pl_module.image_metrics.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()))
+            pl_module.pixel_metrics.to(self.device)
+            pl_module.pixel_metrics.update(torch.squeeze(output["anomaly_maps"]), torch.squeeze(output["mask"].int()))
+        if "semantic_mask" in output and "anomaly_maps" in output:
+            pl_module.semantic_pixel_metrics.to(self.device)
+            pl_module.semantic_pixel_metrics.update(torch.squeeze(output["anomaly_maps"]), output["semantic_mask"])
 
     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
 
     @staticmethod
     def _log_metrics(pl_module: AnomalyModule) -> None:
         """Log computed performance metrics."""
-        if pl_module.pixel_metrics._update_called:  # noqa: SLF001
-            pl_module.log_dict(pl_module.pixel_metrics, prog_bar=True)
-            pl_module.log_dict(pl_module.image_metrics, prog_bar=False)
-        else:
-            pl_module.log_dict(pl_module.image_metrics, prog_bar=True)
+        pl_module.log_dict(pl_module.image_metrics, prog_bar=True)
+        if pl_module.pixel_metrics.update_called:
+            pl_module.log_dict(pl_module.pixel_metrics, prog_bar=False)
+        if pl_module.semantic_pixel_metrics.update_called:
+            pl_module.log_dict(pl_module.semantic_pixel_metrics, prog_bar=False)

src/anomalib/callbacks/normalization/min_max_normalization.py

@@ -39,7 +39,7 @@ def on_test_start(self, trainer: Trainer, pl_module: AnomalyModule) -> None:
         """Call when the test begins."""
         del trainer  # `trainer` variable is not used.
 
-        for metric in (pl_module.image_metrics, pl_module.pixel_metrics):
+        for metric in (pl_module.image_metrics, pl_module.pixel_metrics, pl_module.semantic_pixel_metrics):
             if metric is not None:
                 metric.set_threshold(0.5)
 

src/anomalib/cli/cli.py

@@ -141,8 +141,17 @@ def add_arguments_to_parser(self, parser: ArgumentParser) -> None:
 
         parser.add_function_arguments(get_normalization_callback, "normalization")
         parser.add_argument("--task", type=TaskType | str, 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.image",
+            type=list[str] | str | dict[str, dict[str, Any]] | None,
+            default=["F1Score", "AUROC"],
+        )
+        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 | str, 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 not in ("export", "predict"):

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 .utils import LabelName
 from .video import Avenue, ShanghaiTech, UCSDped, VideoDataFormat
@@ -63,6 +63,7 @@ def get_datamodule(config: DictConfig | ListConfig) -> AnomalibDataModule:
     "Kolektor",
     "MVTec",
     "MVTec3D",
+    "MVTecLoco",
     "Avenue",
     "UCSDped",
     "ShanghaiTech",

src/anomalib/data/base/datamodule.py

@@ -28,7 +28,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` exists,
+    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.
@@ -42,6 +44,12 @@ 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 "semantic_mask" in elem:
+            # semantic masks have a variable number of channels, so we collate them as a list
+            out_dict["semantic_mask"] = [item.pop("semantic_mask") for item in batch]
+        if "mask_path" in elem and isinstance(elem["mask_path"], list):
+            # collate mask paths as list
+            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
@@ -213,6 +221,12 @@ 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
+            assert hasattr(
+                self,
+                "val_data",
+            ), f"FROM_DIR is not supported for {self.__class__.__name__} which does not assign val_data in _setup."
         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"]