Fix get_spatial_ndim for 2D post-EnsureChannelFirst without regressing 3D

The previous commit removed no_channel from all call sites, which broke
3D no-channel tensors (spatial_ndim clamped to ndim-1=2 instead of 3).

Targeted fix: only remove no_channel from get_spatial_ndim, since it runs
after EnsureChannelFirst has already added a channel dim. The constructor
and affine setter keep no_channel to correctly handle pre-channel tensors.

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

Author: aymuos15

monai/data/meta_tensor.py

@@ -23,25 +23,34 @@
 import monai
 from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor
 from monai.data.meta_obj import _DEFAULT_SPATIAL_NDIM, MetaObj, get_track_meta
-from monai.data.utils import affine_to_spacing, decollate_batch, list_data_collate, remove_extra_metadata
+from monai.data.utils import affine_to_spacing, decollate_batch, is_no_channel, list_data_collate, remove_extra_metadata
 from monai.utils import look_up_option
 from monai.utils.enums import LazyAttr, MetaKeys, PostFix, SpaceKeys
 from monai.utils.type_conversion import convert_data_type, convert_to_dst_type, convert_to_numpy, convert_to_tensor
 
 __all__ = ["MetaTensor", "get_spatial_ndim"]
 
 
-def _normalize_spatial_ndim(spatial_ndim: int, tensor_ndim: int) -> int:
+def _normalize_spatial_ndim(spatial_ndim: int, tensor_ndim: int, no_channel: bool = False) -> int:
     """Clamp spatial dims to a valid range for the current tensor shape."""
-    limit = max(int(tensor_ndim) - 1, 1)
+    limit = max(int(tensor_ndim), 1) if no_channel else max(int(tensor_ndim) - 1, 1)
     return max(1, min(int(spatial_ndim), limit))
 
 
+def _has_explicit_no_channel(meta: Mapping | None) -> bool:
+    return (
+        isinstance(meta, Mapping)
+        and MetaKeys.ORIGINAL_CHANNEL_DIM in meta
+        and is_no_channel(meta[MetaKeys.ORIGINAL_CHANNEL_DIM])
+    )
+
+
 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``.
+    ``img.ndim - 1``.  Always assumes channel-first (``no_channel=False``)
+    because callers run after ``EnsureChannelFirst`` has already added one.
     """
     if isinstance(img, MetaTensor):
         return _normalize_spatial_ndim(img.spatial_ndim, img.ndim)
@@ -192,10 +201,11 @@ def __init__(
             self.affine = self.get_default_affine()
         # Initialize spatial_ndim from affine matrix (source of truth), clamped by tensor shape.
         # This cached value is kept in sync via the affine setter for hot-path performance.
+        no_channel = _has_explicit_no_channel(self.meta)
         if spatial_ndim is not None:
-            self.spatial_ndim = _normalize_spatial_ndim(spatial_ndim, self.ndim)
+            self.spatial_ndim = _normalize_spatial_ndim(spatial_ndim, self.ndim, no_channel=no_channel)
         elif self.affine.ndim == 2:
-            self.spatial_ndim = _normalize_spatial_ndim(self.affine.shape[-1] - 1, self.ndim)
+            self.spatial_ndim = _normalize_spatial_ndim(self.affine.shape[-1] - 1, self.ndim, no_channel=no_channel)
 
         # applied_operations
         if applied_operations is not None:
@@ -518,7 +528,8 @@ 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 from affine (source of truth)
-            self.spatial_ndim = _normalize_spatial_ndim(a.shape[-1] - 1, self.ndim)
+            no_channel = _has_explicit_no_channel(self.meta)
+            self.spatial_ndim = _normalize_spatial_ndim(a.shape[-1] - 1, self.ndim, no_channel=no_channel)
 
     @property
     def spatial_ndim(self) -> int: