fix(transforms): guard None spatial_size in spatial_resample (#9068)

Signed-off-by: FinalSunFlower <auroral.sunflower@gmail.com>

Author: FinalSunFlower

monai/transforms/spatial/functional.py

@@ -156,7 +156,7 @@ def spatial_resample(
     elif spatial_size is None and spatial_rank > 1:  # auto spatial size
         spatial_size, _ = compute_shape_offset(in_spatial_size, src_affine, dst_affine)  # type: ignore
     spatial_size = torch.tensor(
-        fall_back_tuple(ensure_tuple(spatial_size)[:spatial_rank], in_spatial_size, lambda x: x >= 0)
+        fall_back_tuple(ensure_tuple(spatial_size)[:spatial_rank], in_spatial_size, lambda x: x is not None and x >= 0)
     )
     extra_info = {
         "dtype": str(dtype_pt)[6:],  # remove "torch": torch.float32 -> float32

tests/transforms/test_spatial_resample.py

@@ -23,7 +23,7 @@
 from monai.data.utils import to_affine_nd
 from monai.transforms import SpatialResample
 from monai.utils import optional_import
-from tests.lazy_transforms_utils import test_resampler_lazy
+from tests.lazy_transforms_utils import test_resampler_lazy as check_resampler_lazy
 from tests.test_utils import TEST_DEVICES, TEST_NDARRAYS_ALL, assert_allclose, dict_product
 
 TESTS = []
@@ -148,7 +148,7 @@ def test_flips(self, img, device, data_param, expected_output):
             assert_allclose(out, expected_output, rtol=1e-2, atol=1e-2)
             assert_allclose(to_affine_nd(len(out.shape) - 1, out.affine), call_param["dst_affine"])
 
-            test_resampler_lazy(resampler, out, init_param=None, call_param=call_param)
+            check_resampler_lazy(resampler, out, init_param=None, call_param=call_param)
 
     @parameterized.expand(TEST_4_5_D)
     def test_4d_5d(self, new_shape, tile, device, dtype, expected_data):
@@ -165,7 +165,7 @@ def test_4d_5d(self, new_shape, tile, device, dtype, expected_data):
         assert_allclose(out, expected_data[None], rtol=1e-2, atol=1e-2)
         assert_allclose(out.affine, dst.to(torch.float32), rtol=1e-2, atol=1e-2)
 
-        test_resampler_lazy(resampler, out, init_param, call_param)
+        check_resampler_lazy(resampler, out, init_param, call_param)
 
     @parameterized.expand(TEST_DEVICES)
     def test_ill_affine(self, device):
@@ -199,7 +199,7 @@ def test_input_torch(self, new_shape, tile, device, dtype, expected_data, track_
         out = resampler(**call_param)
         assert_allclose(out, expected_data[None], rtol=1e-2, atol=1e-2)
 
-        test_resampler_lazy(resampler, out, init_param, call_param)
+        check_resampler_lazy(resampler, out, init_param, call_param)
 
         if track_meta:
             self.assertIsInstance(out, MetaTensor)
@@ -230,6 +230,22 @@ def test_unchange(self):
             assert_allclose(result, img, type_test=False)
         set_track_meta(True)
 
+    def test_none_spatial_size_rank_one(self):
+        img = MetaTensor(torch.randn(1, 8))
+        result = SpatialResample()(img, spatial_size=None)
+
+        self.assertEqual(result.shape, img.shape)
+        self.assertIsInstance(result, MetaTensor)
+        self.assertTrue(torch.isfinite(result).all())
+
+    def test_partial_none_spatial_size(self):
+        img = MetaTensor(torch.randn(1, 3, 6, 7))
+        result = SpatialResample()(img, spatial_size=(None, 4, 5))
+
+        self.assertEqual(result.shape, (1, 3, 4, 5))
+        self.assertIsInstance(result, MetaTensor)
+        self.assertTrue(torch.isfinite(result).all())
+
 
 if __name__ == "__main__":
     unittest.main()