Fix spatial_ndim drift for sliced MetaTensor 2D paths

Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>

Author: aymuos15

monai/data/meta_tensor.py

@@ -31,14 +31,25 @@
 __all__ = ["MetaTensor", "get_spatial_ndim"]
 
 
+def _normalize_spatial_ndim(spatial_ndim: int, tensor_ndim: int) -> int:
+    """Clamp spatial dims to a valid range for the current tensor shape."""
+    return max(1, min(int(spatial_ndim), max(int(tensor_ndim) - 1, 1)))
+
+
 def get_spatial_ndim(img: NdarrayOrTensor) -> int:
     """Return the number of spatial dimensions assuming channel-first layout.
 
     Uses ``MetaTensor.spatial_ndim`` when available, otherwise falls back to
     ``img.ndim - 1``.
     """
     if isinstance(img, MetaTensor):
-        return img.spatial_ndim
+        inferred = _normalize_spatial_ndim(img.spatial_ndim, img.ndim)
+        shape_spatial = max(img.ndim - 1, 1)
+        # For non-batched tensors, preserve explicit higher-rank shape information
+        # (e.g., invalid 4D spatial inputs should still be reported as rank 4).
+        if not img.is_batch and shape_spatial > inferred:
+            return shape_spatial
+        return inferred
     return img.ndim - 1
 
 
@@ -175,9 +186,9 @@ def __init__(
             self.affine = self.get_default_affine()
         # derive spatial_ndim from affine, clamped by tensor shape
         if spatial_ndim is not None:
-            self.spatial_ndim = spatial_ndim
+            self.spatial_ndim = _normalize_spatial_ndim(spatial_ndim, self.ndim)
         elif self.affine.ndim == 2:
-            self.spatial_ndim = min(self.affine.shape[-1] - 1, max(self.ndim - 1, 1))
+            self.spatial_ndim = _normalize_spatial_ndim(self.affine.shape[-1] - 1, self.ndim)
 
         # applied_operations
         if applied_operations is not None:
@@ -243,6 +254,8 @@ def update_meta(rets: Sequence, func, args, kwargs) -> Sequence:
                 #     raise NotImplementedError("torch.cat is not implemented for batch of MetaTensors.")
                 if is_batch:
                     ret = MetaTensor._handle_batched(ret, idx, metas, func, args, kwargs)
+                if func == torch.Tensor.__getitem__:
+                    ret.spatial_ndim = _normalize_spatial_ndim(ret.spatial_ndim, ret.ndim)
             out.append(ret)
         # if the input was a tuple, then return it as a tuple
         return tuple(out) if isinstance(rets, tuple) else out
@@ -492,7 +505,7 @@ def affine(self, d: NdarrayTensor) -> None:
         a = torch.as_tensor(d, device=torch.device("cpu"), dtype=torch.float64)
         self.meta[MetaKeys.AFFINE] = a
         if a.ndim == 2:  # non-batched: sync spatial_ndim
-            self.spatial_ndim = a.shape[-1] - 1
+            self.spatial_ndim = _normalize_spatial_ndim(a.shape[-1] - 1, self.ndim)
 
     @property
     def spatial_ndim(self) -> int:

monai/data/utils.py

@@ -432,7 +432,7 @@ def collate_meta_tensor_fn(batch, *, collate_fn_map=None):
     collated.meta = default_collate(meta_dicts)
     collated.applied_operations = [i.applied_operations or TraceKeys.NONE for i in batch]
     collated.is_batch = True
-    collated.spatial_ndim = getattr(batch[0], "spatial_ndim", 3)  # assumes uniform spatial_ndim
+    collated.spatial_ndim = min(getattr(batch[0], "spatial_ndim", 3), max(collated.ndim - 1, 1))
     return collated
 
 

tests/data/meta_tensor/test_spatial_ndim.py

@@ -19,9 +19,9 @@
 import torch
 from parameterized import parameterized
 
-from monai.data import MetaTensor
+from monai.data import MetaTensor, get_spatial_ndim
 from monai.data.utils import collate_meta_tensor_fn, decollate_batch
-from monai.transforms import Affine, RandAffine, Resize, Rotate, SqueezeDim
+from monai.transforms import Affine, LabelToContour, RandAffine, RandZoom, Resize, Rotate, SqueezeDim
 from monai.transforms.utility.array import SplitDim
 from monai.utils import optional_import
 
@@ -121,6 +121,30 @@ def test_lazy_apply_pending_2d(self):
         self.assertIsInstance(result, MetaTensor)
         self.assertEqual(len(applied), 1)
 
+    def test_batch_slice_clamps_spatial_ndim(self):
+        t = MetaTensor(torch.randn(10, 6, 5, 7), affine=torch.eye(4))
+        self.assertEqual(t.spatial_ndim, 3)
+        sliced = t[0]
+        self.assertEqual(sliced.shape, (6, 5, 7))
+        self.assertEqual(sliced.spatial_ndim, 2)
+        self.assertEqual(get_spatial_ndim(sliced), 2)
+
+    def test_label_to_contour_batch_slice_2d(self):
+        t = MetaTensor(torch.randint(0, 2, (10, 6, 5, 7)).float(), affine=torch.eye(4))
+        sliced = t[0]
+        out = LabelToContour()(sliced)
+        self.assertEqual(out.shape, sliced.shape)
+
+    def test_rand_zoom_batch_slice_2d(self):
+        t = MetaTensor(torch.randn(10, 1, 64, 64), affine=torch.eye(4))
+        sliced = t[0]
+        zoom = RandZoom(prob=1.0, min_zoom=0.6, max_zoom=1.2)
+        zoom.set_random_state(seed=0)
+        zoom.randomize(sliced)
+        self.assertEqual(len(zoom._zoom), 2)
+        out = zoom(sliced)
+        self.assertEqual(out.ndim, sliced.ndim)
+
     @skipUnless(has_einops, "Requires einops")
     def test_einops_rearrange_then_resize(self):
         """Reproduce the exact #6397 bug: einops.rearrange -> Resize."""