Dataclasses and post-processing refactor (open-edge-platform#2098)

* use dataclass for model in- and outputs

* split dataclass in image and video

* use dataclass in torch inferencer

* use dataclass in openvino inferencer

* add post_processor class

* remove default metrics from CLI

* export post processing

* add post processor to patchcore

* use named tuple for inference outputs

* validate and format inputs of PredictBatch

* update torch inference

* remove base inferencer inheritance

* update openvino inference

* fix visualization

* PredictBatch -> Batch

* post processor as callback

* use callback methods to apply post processing

* temporary fix for visualization

* add DatasetItem class

* fix pred_score shape and add __len__

* make batch iterable

* add in place replace method

* use dataset items in inference

* dataset_items -> items

* use namedtuple as torch model outputs

* formatting

* split dataclasses into input/output and image/video

* merge input and output classes

* use init_subclass for attribute checking

* add descriptor class for validation

* improve error handling

* DataClassDescriptor -> FieldDescriptor

* add is_optional method

* add input validation for torch image and batch

* use image and video dataclasses in library

* add more validation

* add validation

* make postprocessor configurable from engine

* fix post processing logic

* fix data tests

* remove detection task type

* fix more tests

* use separate normalization stats for image and pixel preds

* add sensitivity parameters to one class pp

* fix utils tests

* fix utils tests

* remove metric serialization test

* remove normalization and thresholding args

* set default post processor in base model

* remove manual threshold test

* fix remaining unit tests

* add post_processor to CLI args

* remove old post processing callbacks

* remove comment

* remove references to old normalization and thresholding callbacks

* remove reshape in openvino inferencer

* export lightning model directly

* make collate accessible from dataset

* fix tools integration tests

* add update method to dataclasses

* allow missing pred_score or anomaly_map in post processor

* fix exportable centercrop conversion

* fix model tests

* test all models

* fix efficient_ad

* post processor as model arg

* disable rkde tests

* fix winclip export

* add copyright notice

* add validation for numpy anomaly map

* fix getting started notebook

* remove hardcoded path

* update dataset notebooks

* update model notebooks

* fix logging notebooks

* fix model notebook

Author: djdameln

notebooks/000_getting_started/001_getting_started.ipynb

@@ -20,5 +20,4 @@ class TaskType(str, Enum):
     """Task type used when generating predictions on the dataset."""
 
     CLASSIFICATION = "classification"
-    DETECTION = "detection"
     SEGMENTATION = "segmentation"

notebooks/100_datamodules/101_btech.ipynb

@@ -4,6 +4,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
+from dataclasses import asdict
 from enum import Enum
 from typing import Any
 
@@ -12,6 +13,7 @@
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 
 from anomalib import TaskType
+from anomalib.dataclasses import Batch
 from anomalib.metrics import AnomalibMetricCollection, create_metric_collection
 from anomalib.models import AnomalyModule
 
@@ -96,7 +98,6 @@ def setup(
                         pl_module.pixel_metrics.add_metrics(new_metrics[name])
             else:
                 pl_module.pixel_metrics = create_metric_collection(pixel_metric_names, "pixel_")
-            self._set_threshold(pl_module)
 
     @staticmethod
     def on_validation_epoch_start(trainer: Trainer, pl_module: AnomalyModule) -> None:
@@ -117,13 +118,12 @@ 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(self, trainer: Trainer, pl_module: AnomalyModule) -> None:
         del trainer  # Unused argument.
 
-        self._set_threshold(pl_module)
         self._log_metrics(pl_module)
 
     @staticmethod
@@ -145,35 +145,32 @@ 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(self, trainer: Trainer, pl_module: AnomalyModule) -> None:
         del trainer  # Unused argument.
 
         self._log_metrics(pl_module)
 
-    @staticmethod
-    def _set_threshold(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 _update_metrics(
         self,
         image_metric: AnomalibMetricCollection,
         pixel_metric: AnomalibMetricCollection,
         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, Batch):
+            output = output.__class__(**self._outputs_to_device(asdict(output)))
         elif isinstance(output, torch.Tensor):
             output = output.to(self.device)
         return output