fix video mask

Signed-off-by: Ashwin Vaidya <[email protected]>

Author: ashwinvaidya17

src/anomalib/data/validators/torch/video.py

@@ -600,17 +600,18 @@ def validate_gt_mask(mask: torch.Tensor | None) -> Mask | None:
         if mask is None:
             return None
         if not isinstance(mask, torch.Tensor):
-            msg = f"Masks must be a torch.Tensor, got {type(mask)}."
+            msg = f"Ground truth mask must be a torch.Tensor, got {type(mask)}."
             raise TypeError(msg)
-        if mask.ndim not in {3, 4, 5}:
-            msg = f"Masks must have shape [B, H, W], [B, T, H, W] or [B, T, 1, H, W], got shape {mask.shape}."
+        if mask.ndim not in {2, 3, 4}:
+            msg = f"Ground truth mask must have shape [H, W] or [N, H, W] or [N, 1, H, W] got shape {mask.shape}."
             raise ValueError(msg)
-        if mask.ndim == 5:
-            if mask.shape[2] != 1:
-                msg = f"Masks must have 1 channel, got {mask.shape[2]}."
+        if mask.ndim == 2:
+            mask = mask.unsqueeze(0)
+        if mask.ndim == 4:
+            if mask.shape[1] != 1:
+                msg = f"Ground truth mask must have 1 channel, got {mask.shape[1]}."
                 raise ValueError(msg)
-            mask = mask.squeeze(2)
-
+            mask = mask.squeeze(1)
         return Mask(mask, dtype=torch.bool)
 
     @staticmethod

src/anomalib/models/video/ai_vad/lightning_model.py

@@ -7,7 +7,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
-from dataclasses import replace
 from typing import Any
 
 from lightning.pytorch.utilities.types import STEP_OUTPUT
@@ -146,13 +145,7 @@ def validation_step(self, batch: VideoBatch, *args, **kwargs) -> STEP_OUTPUT:
         del args, kwargs  # Unused arguments.
 
         predictions = self.model(batch.image)
-
-        return replace(
-            batch,
-            pred_score=predictions.pred_score,
-            anomaly_map=predictions.anomaly_map,
-            pred_mask=predictions.pred_mask,
-        )
+        return batch.update(pred_score=predictions.pred_score, anomaly_map=predictions.anomaly_map)
 
     @property
     def trainer_arguments(self) -> dict[str, Any]: