Fix 2D inverse transform shape mismatch (4x4 vs 3x3)

Use affine matrix rank instead of spatial_ndim in Rotate, Affine, and
RandAffine inverse methods to avoid RuntimeError when spatial_ndim=2
but the stored affine is 4x4. Add regression tests for all three.

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

Author: aymuos15

monai/transforms/spatial/array.py

@@ -1036,7 +1036,7 @@ def inverse_transform(self, data: torch.Tensor, transform) -> torch.Tensor:
         out = convert_to_dst_type(out, dst=data, dtype=out.dtype)[0]
         if isinstance(out, MetaTensor):
             affine = convert_to_tensor(out.peek_pending_affine(), track_meta=False)
-            mat = to_affine_nd(out.spatial_ndim, transform_t)
+            mat = to_affine_nd(len(affine) - 1, transform_t)
             out.affine @= convert_to_dst_type(mat, affine)[0]
         return out
 
@@ -2353,7 +2353,7 @@ def inverse(self, data: torch.Tensor) -> torch.Tensor:
         out.meta = data.meta  # type: ignore
         affine = convert_data_type(out.peek_pending_affine(), torch.Tensor)[0]
         xform, *_ = convert_to_dst_type(
-            Affine.compute_w_affine(out.spatial_ndim, inv_affine, data.shape[1:], orig_size), affine
+            Affine.compute_w_affine(len(affine) - 1, inv_affine, data.shape[1:], orig_size), affine
         )
         out.affine @= xform
         return out
@@ -2622,7 +2622,7 @@ def inverse(self, data: torch.Tensor) -> torch.Tensor:
         out.meta = data.meta  # type: ignore
         affine = convert_data_type(out.peek_pending_affine(), torch.Tensor)[0]
         xform, *_ = convert_to_dst_type(
-            Affine.compute_w_affine(out.spatial_ndim, inv_affine, data.shape[1:], orig_size), affine
+            Affine.compute_w_affine(len(affine) - 1, inv_affine, data.shape[1:], orig_size), affine
         )
         out.affine @= xform
         return out

tests/data/meta_tensor/test_spatial_ndim.py

@@ -15,12 +15,13 @@
 from copy import deepcopy
 from unittest import skipUnless
 
+import numpy as np
 import torch
 from parameterized import parameterized
 
 from monai.data import MetaTensor
 from monai.data.utils import collate_meta_tensor_fn, decollate_batch
-from monai.transforms import Resize, SqueezeDim
+from monai.transforms import Affine, RandAffine, Resize, Rotate, SqueezeDim
 from monai.transforms.utility.array import SplitDim
 from monai.utils import optional_import
 
@@ -134,6 +135,37 @@ def test_einops_rearrange_then_resize(self):
         out = Resize(spatial_size=(32, 32, 3), mode="trilinear", align_corners=True)(x_)
         self.assertEqual(out.shape[-3:], (32, 32, 3))
 
+    def test_affine_inverse_2d_metatensor(self):
+        """Affine.inverse on 2D data: 4x4 affine with spatial_ndim=2."""
+        img = MetaTensor(torch.randn(1, 32, 32), affine=torch.eye(4))
+        self.assertEqual(img.spatial_ndim, 2)
+        xform = Affine(rotate_params=(np.pi / 6,), padding_mode="zeros", image_only=True)
+        result = xform(img)
+        inv = xform.inverse(result)
+        self.assertEqual(inv.shape, img.shape)
+        self.assertEqual(len(inv.applied_operations), 0)
+
+    def test_rotate_inverse_2d_metatensor(self):
+        """Rotate.inverse on 2D data: 4x4 affine with spatial_ndim=2."""
+        img = MetaTensor(torch.randn(1, 32, 32), affine=torch.eye(4))
+        self.assertEqual(img.spatial_ndim, 2)
+        xform = Rotate(angle=(np.pi / 4,), padding_mode="zeros")
+        result = xform(img)
+        inv = xform.inverse(result)
+        self.assertEqual(inv.shape, img.shape)
+        self.assertEqual(len(inv.applied_operations), 0)
+
+    def test_rand_affine_inverse_2d_metatensor(self):
+        """RandAffine.inverse on 2D data: 4x4 affine with spatial_ndim=2."""
+        img = MetaTensor(torch.randn(1, 32, 32), affine=torch.eye(4))
+        self.assertEqual(img.spatial_ndim, 2)
+        xform = RandAffine(prob=1.0, rotate_range=(np.pi / 6,), padding_mode="zeros")
+        xform.set_random_state(seed=42)
+        result = xform(img)
+        inv = xform.inverse(result)
+        self.assertEqual(inv.shape, img.shape)
+        self.assertEqual(len(inv.applied_operations), 0)
+
 
 if __name__ == "__main__":
     unittest.main()