Dataclasses and post-processing refactor

src/anomalib/callbacks/metrics.py

@@ -15,6 +15,9 @@
 from anomalib import TaskType
 from anomalib.metrics import AnomalibMetricCollection, create_metric_collection
 from anomalib.models import AnomalyModule
+from anomalib.dataclasses import BatchItem
+
+from dataclasses import asdict
 
 logger = logging.getLogger(__name__)
 
@@ -121,7 +124,7 @@ def on_validation_batch_end(
         del trainer, batch, batch_idx, dataloader_idx  # Unused arguments.
 
         if outputs is not None:
-            self._outputs_to_device(outputs)
+            outputs = self._outputs_to_device(outputs)
             self._update_metrics(pl_module.image_metrics, pl_module.pixel_metrics, outputs)
 
     def on_validation_epoch_end(
@@ -156,7 +159,7 @@ def on_test_batch_end(
         del trainer, batch, batch_idx, dataloader_idx  # Unused arguments.
 
         if outputs is not None:
-            self._outputs_to_device(outputs)
+            outputs = self._outputs_to_device(outputs)
             self._update_metrics(pl_module.image_metrics, pl_module.pixel_metrics, outputs)
 
     def on_test_epoch_end(
@@ -179,15 +182,17 @@ def _update_metrics(
         output: STEP_OUTPUT,
     ) -> None:
         image_metric.to(self.device)
-        image_metric.update(output["pred_scores"], output["label"].int())
-        if "mask" in output and "anomaly_maps" in output:
+        image_metric.update(output.pred_score, output.gt_label.int())
+        if output.gt_mask is not None and output.anomaly_map is not None:
             pixel_metric.to(self.device)
-            pixel_metric.update(torch.squeeze(output["anomaly_maps"]), torch.squeeze(output["mask"].int()))
+            pixel_metric.update(torch.squeeze(output.anomaly_map), torch.squeeze(output.gt_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, BatchItem):
+            output = output.__class__(**self._outputs_to_device(asdict(output)))
         elif isinstance(output, torch.Tensor):
             output = output.to(self.device)
         return output

src/anomalib/callbacks/normalization/min_max_normalization.py

@@ -55,12 +55,12 @@ def on_validation_batch_end(
         """Call when the validation batch ends, update the min and max observed values."""
         del trainer, batch, batch_idx, dataloader_idx  # These variables are not used.
 
-        if "anomaly_maps" in outputs:
-            pl_module.normalization_metrics(outputs["anomaly_maps"])
-        elif "box_scores" in outputs:
-            pl_module.normalization_metrics(torch.cat(outputs["box_scores"]))
-        elif "pred_scores" in outputs:
-            pl_module.normalization_metrics(outputs["pred_scores"])
+        if outputs.anomaly_map is not None:
+            pl_module.normalization_metrics(outputs.anomaly_map)
+        elif outputs.box_scores is not None:
+            pl_module.normalization_metrics(torch.cat(outputs.box_scores))
+        elif outputs.pred_score is not None:
+            pl_module.normalization_metrics(outputs.pred_score)
         else:
             msg = "No values found for normalization, provide anomaly maps, bbox scores, or image scores"
             raise ValueError(msg)
@@ -99,11 +99,11 @@ def _normalize_batch(outputs: Any, pl_module: AnomalyModule) -> None:  # noqa: A
         image_threshold = pl_module.image_threshold.value.cpu()
         pixel_threshold = pl_module.pixel_threshold.value.cpu()
         stats = pl_module.normalization_metrics.cpu()
-        if "pred_scores" in outputs:
-            outputs["pred_scores"] = normalize(outputs["pred_scores"], image_threshold, stats.min, stats.max)
-        if "anomaly_maps" in outputs:
-            outputs["anomaly_maps"] = normalize(outputs["anomaly_maps"], pixel_threshold, stats.min, stats.max)
-        if "box_scores" in outputs:
-            outputs["box_scores"] = [
-                normalize(scores, pixel_threshold, stats.min, stats.max) for scores in outputs["box_scores"]
+        if outputs.pred_score is not None:
+            outputs.pred_score = normalize(outputs.pred_score, image_threshold, stats.min, stats.max)
+        if outputs.anomaly_map is not None:
+            outputs.anomaly_map = normalize(outputs.anomaly_map, pixel_threshold, stats.min, stats.max)
+        if outputs.box_scores is not None:
+            outputs.box_scores = [
+                normalize(scores, pixel_threshold, stats.min, stats.max) for scores in outputs.box_scores
             ]

src/anomalib/callbacks/post_processor.py

@@ -13,6 +13,7 @@
 
 from anomalib.data.utils import boxes_to_anomaly_maps, boxes_to_masks, masks_to_boxes
 from anomalib.models import AnomalyModule
+from anomalib.dataclasses import BatchItem
 
 
 class _PostProcessorCallback(Callback):
@@ -28,7 +29,7 @@ def on_validation_batch_end(
         self,
         trainer: Trainer,
         pl_module: AnomalyModule,
-        outputs: STEP_OUTPUT | None,
+        outputs: BatchItem,
         batch: Any,  # noqa: ANN401
         batch_idx: int,
         dataloader_idx: int = 0,
@@ -42,7 +43,7 @@ def on_test_batch_end(
         self,
         trainer: Trainer,
         pl_module: AnomalyModule,
-        outputs: STEP_OUTPUT | None,
+        outputs: BatchItem,
         batch: Any,  # noqa: ANN401
         batch_idx: int,
         dataloader_idx: int = 0,
@@ -56,7 +57,7 @@ def on_predict_batch_end(
         self,
         trainer: Trainer,
         pl_module: AnomalyModule,
-        outputs: Any,  # noqa: ANN401
+        outputs: BatchItem,  # noqa: ANN401
         batch: Any,  # noqa: ANN401
         batch_idx: int,
         dataloader_idx: int = 0,
@@ -67,7 +68,7 @@ def on_predict_batch_end(
             self.post_process(trainer, pl_module, outputs)
 
     def post_process(self, trainer: Trainer, pl_module: AnomalyModule, outputs: STEP_OUTPUT) -> None:
-        if isinstance(outputs, dict):
+        if isinstance(outputs, BatchItem):
             self._post_process(outputs)
             if trainer.predicting or trainer.testing:
                 self._compute_scores_and_labels(pl_module, outputs)
@@ -77,49 +78,49 @@ def _compute_scores_and_labels(
         pl_module: AnomalyModule,
         outputs: dict[str, Any],
     ) -> None:
-        if "pred_scores" in outputs:
-            outputs["pred_labels"] = outputs["pred_scores"] >= pl_module.image_threshold.value
-        if "anomaly_maps" in outputs:
-            outputs["pred_masks"] = outputs["anomaly_maps"] >= pl_module.pixel_threshold.value
-            if "pred_boxes" not in outputs:
-                outputs["pred_boxes"], outputs["box_scores"] = masks_to_boxes(
-                    outputs["pred_masks"],
-                    outputs["anomaly_maps"],
+        if outputs.pred_score is not None:
+            outputs.pred_label = outputs.pred_score >= pl_module.image_threshold.value
+        if outputs.anomaly_map is not None:
+            outputs.pred_mask = outputs.anomaly_map >= pl_module.pixel_threshold.value
+            if outputs.pred_boxes is None:
+                outputs.pred_boxes, outputs.box_scores = masks_to_boxes(
+                    outputs.pred_mask,
+                    outputs.anomaly_map,
                 )
-                outputs["box_labels"] = [torch.ones(boxes.shape[0]) for boxes in outputs["pred_boxes"]]
+                outputs.box_labels = [torch.ones(boxes.shape[0]) for boxes in outputs.pred_boxes]
         # apply thresholding to boxes
-        if "box_scores" in outputs and "box_labels" not in outputs:
+        if outputs.box_scores is not None and outputs.box_labels is None:
             # apply threshold to assign normal/anomalous label to boxes
-            is_anomalous = [scores > pl_module.pixel_threshold.value for scores in outputs["box_scores"]]
-            outputs["box_labels"] = [labels.int() for labels in is_anomalous]
+            is_anomalous = [scores > pl_module.pixel_threshold.value for scores in outputs.box_scores]
+            outputs.box_labels = [labels.int() for labels in is_anomalous]
 
     @staticmethod
-    def _post_process(outputs: STEP_OUTPUT) -> None:
+    def _post_process(outputs: BatchItem) -> None:
         """Compute labels based on model predictions."""
-        if isinstance(outputs, dict):
-            if "pred_scores" not in outputs and "anomaly_maps" in outputs:
+        if isinstance(outputs, BatchItem):
+            if outputs.pred_score is None and outputs.anomaly_map is not None:
                 # infer image scores from anomaly maps
-                outputs["pred_scores"] = (
-                    outputs["anomaly_maps"]  # noqa: PD011
-                    .reshape(outputs["anomaly_maps"].shape[0], -1)
+                outputs.pred_score = (
+                    outputs.anomaly_map  # noqa: PD011
+                    .reshape(outputs.anomaly_map.shape[0], -1)
                     .max(dim=1)
                     .values
                 )
-            elif "pred_scores" not in outputs and "box_scores" in outputs and "label" in outputs:
+            elif outputs.pred_score is None and outputs.box_score is not None and outputs.gt_label is not None:
                 # infer image score from bbox confidence scores
-                outputs["pred_scores"] = torch.zeros_like(outputs["label"]).float()
-                for idx, (boxes, scores) in enumerate(zip(outputs["pred_boxes"], outputs["box_scores"], strict=True)):
+                outputs.pred_score = torch.zeros_like(outputs.gt_label).float()
+                for idx, (boxes, scores) in enumerate(zip(outputs.pred_boxes, outputs.box_scores, strict=True)):
                     if boxes.numel():
-                        outputs["pred_scores"][idx] = scores.max().item()
+                        outputs.pred_score[idx] = scores.max().item()
 
-            if "pred_boxes" in outputs and "anomaly_maps" not in outputs:
+            if outputs.pred_boxes is not None and outputs.anomaly_map is None:
                 # create anomaly maps from bbox predictions for thresholding and evaluation
-                image_size: tuple[int, int] = outputs["image"].shape[-2:]
-                pred_boxes: torch.Tensor = outputs["pred_boxes"]
-                box_scores: torch.Tensor = outputs["box_scores"]
+                image_size: tuple[int, int] = outputs.image.shape[-2:]
+                pred_boxes: torch.Tensor = outputs.pred_boxes
+                box_scores: torch.Tensor = outputs.box_scores
 
-                outputs["anomaly_maps"] = boxes_to_anomaly_maps(pred_boxes, box_scores, image_size)
+                outputs.anomaly_map = boxes_to_anomaly_maps(pred_boxes, box_scores, image_size)
 
-                if "boxes" in outputs:
-                    true_boxes: list[torch.Tensor] = outputs["boxes"]
-                    outputs["mask"] = boxes_to_masks(true_boxes, image_size)
+                if outputs.gt_boxes is not None:
+                    true_boxes: list[torch.Tensor] = outputs.gt_boxes
+                    outputs.gt_mask = boxes_to_masks(true_boxes, image_size)

src/anomalib/callbacks/thresholding.py

@@ -14,6 +14,9 @@
 from anomalib.metrics.threshold import BaseThreshold
 from anomalib.models import AnomalyModule
 from anomalib.utils.types import THRESHOLD
+from anomalib.dataclasses import BatchItem
+
+from dataclasses import asdict
 
 
 class _ThresholdCallback(Callback):
@@ -53,7 +56,7 @@ def on_validation_batch_end(
     ) -> None:
         del trainer, batch, batch_idx, dataloader_idx  # Unused arguments.
         if outputs is not None:
-            self._outputs_to_cpu(outputs)
+            outputs = self._outputs_to_cpu(outputs)
             self._update(pl_module, outputs)
 
     def on_validation_epoch_end(self, trainer: Trainer, pl_module: AnomalyModule) -> None:
@@ -178,16 +181,18 @@ def _outputs_to_cpu(self, output: STEP_OUTPUT) -> STEP_OUTPUT | dict[str, Any]:
         if isinstance(output, dict):
             for key, value in output.items():
                 output[key] = self._outputs_to_cpu(value)
+        elif isinstance(output, BatchItem):
+            output = output.__class__(**self._outputs_to_cpu(asdict(output)))
         elif isinstance(output, torch.Tensor):
             output = output.cpu()
         return output
 
     def _update(self, pl_module: AnomalyModule, outputs: STEP_OUTPUT) -> None:
         pl_module.image_threshold.cpu()
-        pl_module.image_threshold.update(outputs["pred_scores"], outputs["label"].int())
-        if "mask" in outputs and "anomaly_maps" in outputs:
+        pl_module.image_threshold.update(outputs.pred_score, outputs.gt_label.int())
+        if outputs.gt_mask is not None and outputs.anomaly_map is not None:
             pl_module.pixel_threshold.cpu()
-            pl_module.pixel_threshold.update(outputs["anomaly_maps"], outputs["mask"].int())
+            pl_module.pixel_threshold.update(outputs.anomaly_map, outputs.gt_mask.int())
 
     def _compute(self, pl_module: AnomalyModule) -> None:
         pl_module.image_threshold.compute()

src/anomalib/data/base/datamodule.py

@@ -17,6 +17,8 @@
 
 from anomalib.data.utils import TestSplitMode, ValSplitMode, random_split, split_by_label
 from anomalib.data.utils.synthetic import SyntheticAnomalyDataset
+from anomalib.dataclasses import BatchItem
+from dataclasses import asdict
 
 if TYPE_CHECKING:
     from pandas import DataFrame
@@ -26,7 +28,7 @@
 logger = logging.getLogger(__name__)
 
 
-def collate_fn(batch: list) -> dict[str, Any]:
+def collate_fn(batch: list[BatchItem]) -> 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.
@@ -37,16 +39,18 @@ def collate_fn(batch: list) -> dict[str, Any]:
     Returns:
         dict[str, Any]: Dictionary containing the collated batch information.
     """
-    elem = batch[0]  # sample an element from the batch to check the type.
+    # convert to list of dicts
+    batch_dict = [asdict(item) for item in batch]
+    elem = batch_dict[0]  # sample an element from the batch to check the type.
     out_dict = {}
-    if isinstance(elem, dict):
-        if "boxes" in elem:
-            # collate boxes as list
-            out_dict["boxes"] = [item.pop("boxes") 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
-    return default_collate(batch)
+    # if isinstance(elem, dict):
+    if "boxes" in elem:
+        # collate boxes as list
+        out_dict["boxes"] = [item.pop("boxes") for item in batch_dict]
+    # collate other data normally
+    out_dict.update({key: default_collate([item[key] for item in batch]) for key in elem if elem[key] is not None})
+    return batch[0].__class__(**out_dict)
+    # return default_collate(batch)
 
 
 class AnomalibDataModule(LightningDataModule, ABC):
@@ -225,6 +229,7 @@ def train_dataloader(self) -> TRAIN_DATALOADERS:
             shuffle=True,
             batch_size=self.train_batch_size,
             num_workers=self.num_workers,
+            collate_fn=collate_fn,
         )
 
     def val_dataloader(self) -> EVAL_DATALOADERS:

src/anomalib/data/base/dataset.py

@@ -17,6 +17,7 @@
 from torchvision.tv_tensors import Mask
 
 from anomalib import TaskType
+from anomalib.dataclasses import ImageBatch
 from anomalib.data.utils import LabelName, masks_to_boxes, read_image, read_mask
 
 _EXPECTED_COLUMNS_CLASSIFICATION = ["image_path", "split"]
@@ -189,7 +190,14 @@ def __getitem__(self, index: int) -> dict[str, str | torch.Tensor]:
             msg = f"Unknown task type: {self.task}"
             raise ValueError(msg)
 
-        return item
+        # return item
+        return ImageBatch(
+            image=item["image"],
+            gt_mask=item["mask"],
+            gt_label=label_index,
+            image_path=image_path,
+            mask_path=mask_path,
+        )
 
     def __add__(self, other_dataset: "AnomalibDataset") -> "AnomalibDataset":
         """Concatenate this dataset with another dataset.

src/anomalib/data/base/video.py

@@ -129,16 +129,16 @@ def _select_targets(self, item: dict[str, Any]) -> dict[str, Any]:
             msg = f"Unknown video target frame: {self.target_frame}"
             raise ValueError(msg)
 
-        if item.get("mask") is not None:
-            item["mask"] = item["mask"][idx, ...]
-        if item.get("boxes") is not None:
-            item["boxes"] = item["boxes"][idx]
-        if item.get("label") is not None:
-            item["label"] = item["label"][idx]
-        if item.get("original_image") is not None:
-            item["original_image"] = item["original_image"][idx]
-        if item.get("frames") is not None:
-            item["frames"] = item["frames"][idx]
+        if item.mask is not None:
+            item.mask = item.mask[idx, ...]
+        if item.gt_boxes is not None:
+            item.gt_boxes = item.gt_boxes[idx]
+        if item.gt_label is not None:
+            item.gt_label = item.gt_label[idx]
+        if item.original_image is not None:
+            item.original_image = item.original_image[idx]
+        if item.frames is not None:
+            item.frames = item.frames[idx]
         return item
 
     def __getitem__(self, index: int) -> dict[str, str | torch.Tensor]:
@@ -154,30 +154,30 @@ def __getitem__(self, index: int) -> dict[str, str | torch.Tensor]:
             msg = "self.indexer must be an instance of ClipsIndexer."
             raise TypeError(msg)
         item = self.indexer.get_item(index)
-        item["image"] = to_dtype_video(video=item["image"], scale=True)
+        item.image = to_dtype_video(video=item.image, scale=True)
         # include the untransformed image for visualization
-        item["original_image"] = item["image"].to(torch.uint8)
+        item.original_image = item.image.to(torch.uint8)
 
         # apply transforms
-        if item.get("mask") is not None:
+        if item.mask is not None:
             if self.transform:
-                item["image"], item["mask"] = self.transform(item["image"], Mask(item["mask"]))
-            item["label"] = torch.Tensor([1 in frame for frame in item["mask"]]).int().squeeze(0)
+                item.image, item.mask = self.transform(item.image, Mask(item.mask))
+            item.gt_label = torch.Tensor([1 in frame for frame in item.mask]).int().squeeze(0)
             if self.task == TaskType.DETECTION:
-                item["boxes"], _ = masks_to_boxes(item["mask"])
-                item["boxes"] = item["boxes"][0] if len(item["boxes"]) == 1 else item["boxes"]
+                item.gt_boxes, _ = masks_to_boxes(item.mask)
+                item.gt_boxes = item.gt_boxes[0] if len(item.gt_boxes) == 1 else item.gt_boxes
         elif self.transform:
-            item["image"] = self.transform(item["image"])
+            item.image = self.transform(item.image)
 
         # squeeze temporal dimensions in case clip length is 1
-        item["image"] = item["image"].squeeze(0)
+        item.image = item.image.squeeze(0)
 
         # include only target frame in gt
         if self.clip_length_in_frames > 1 and self.target_frame != VideoTargetFrame.ALL:
             item = self._select_targets(item)
 
-        if item["mask"] is None:
-            item.pop("mask")
+        # if item.mask is None:
+        #     item.pop("mask")
 
         return item