diff --git a/monai/transforms/spatial/functional.py b/monai/transforms/spatial/functional.py index 967020a3bf..add5bda8bd 100644 --- a/monai/transforms/spatial/functional.py +++ b/monai/transforms/spatial/functional.py @@ -16,6 +16,7 @@ import math import warnings +from collections.abc import Sequence from enum import Enum import numpy as np @@ -106,6 +107,41 @@ def _maybe_new_metatensor(img, dtype=None, device=None): ) +def _cancel_align_corners_scaling( + xform: torch.Tensor, in_spatial_size: Sequence[int], out_spatial_size: Sequence[int] +) -> torch.Tensor: + """ + Pre-compensate a voxel-space affine for the scalings applied again downstream when + ``align_corners=True`` and the compiled ``Resample`` path is used. + + ``xform`` is already expressed in voxel space. However, ``Affine`` passes it through ``AffineGrid``, + which applies a ``d_out / (d_out - 1)`` scaling per spatial dimension, and the compiled ``Resample`` + applies a ``(d_in - 1) / d_in`` scaling before ``grid_pull``. This function returns + ``T @ S_inv @ xform @ D_inv``, where ``S_inv`` and ``D_inv`` cancel those two scalings and ``T`` is a + translation by -0.5 voxel that re-centres the result, so that ``grid_pull`` receives the original + ``xform`` unchanged. + + Dimensions are clamped to a minimum of 2 to avoid division by zero for singleton axes. + + Args: + xform: homogeneous affine of shape ``(rank + 1, rank + 1)`` in voxel space. + in_spatial_size: spatial shape of the input image. + out_spatial_size: spatial shape of the resampled output. + + Returns: + the compensated homogeneous affine, same shape, dtype and device as ``xform``. + """ + rank = len(xform) - 1 + eye = torch.eye(rank + 1, dtype=xform.dtype, device=xform.device) + s_inv, d_inv, translate = eye.clone(), eye.clone(), eye.clone() + for i in range(rank): + d_in, d_out = max(2, int(in_spatial_size[i])), max(2, int(out_spatial_size[i])) + s_inv[i, i] = d_in / (d_in - 1) + d_inv[i, i] = (d_out - 1) / d_out + translate[i, -1] = -0.5 + return translate @ s_inv @ xform @ d_inv + + def spatial_resample( img, dst_affine, spatial_size, mode, padding_mode, align_corners, dtype_pt, lazy, transform_info ) -> torch.Tensor: @@ -216,6 +252,8 @@ def spatial_resample( if isinstance(mode, int) or _use_compiled: dst_xform = create_translate(spatial_rank, [float(d - 1) / 2 for d in spatial_size]) xform = xform @ convert_to_dst_type(dst_xform, xform)[0] + if align_corners and not isinstance(mode, int): + xform = _cancel_align_corners_scaling(xform, in_sp_size, spatial_size) affine_xform = monai.transforms.Affine( affine=xform, spatial_size=spatial_size, diff --git a/tests/transforms/test_affine.py b/tests/transforms/test_affine.py index 8a29e31de8..ec99c65dfd 100644 --- a/tests/transforms/test_affine.py +++ b/tests/transforms/test_affine.py @@ -19,10 +19,10 @@ from parameterized import parameterized from monai.data import MetaTensor, set_track_meta -from monai.transforms import Affine, Resize +from monai.transforms import Affine, RandRotate, Resize from monai.transforms.lazy.functional import apply_pending from monai.utils import optional_import -from tests.lazy_transforms_utils import test_resampler_lazy +from tests.lazy_transforms_utils import get_apply_param, test_resampler_lazy from tests.test_utils import TEST_NDARRAYS_ALL, assert_allclose, test_local_inversion TESTS = [] @@ -291,6 +291,26 @@ def method_3(im, ac): ref = Resize(align_corners=ac, spatial_size=(sp_size, sp_size), mode="bilinear")(im) assert_allclose(out, ref, rtol=1e-4, atol=1e-4, type_test=False) + def test_lazy_rotate_align_corners_compiled(self): + """lazy resampling with the compiled backend and align_corners=True must match eager.""" + init = dict( + range_x=np.pi, + prob=1.0, + keep_size=False, + mode="nearest", + padding_mode="zeros", + dtype=np.float64, + align_corners=True, + ) + img = MetaTensor(torch.arange(64 * 48, dtype=torch.float32).reshape(1, 64, 48)) + xform = RandRotate(**init) + xform.set_random_state(seed=243) + eager = xform(img) + xform.set_random_state(seed=243) + xform.lazy = True + lazy = apply_pending(xform(img), overrides=get_apply_param(init, {"img": img}))[0] + assert_allclose(lazy, eager, rtol=1e-3, atol=1e-3) + if __name__ == "__main__": unittest.main() diff --git a/tests/transforms/test_spacing.py b/tests/transforms/test_spacing.py index b8f838c4ac..40da7ebd1e 100644 --- a/tests/transforms/test_spacing.py +++ b/tests/transforms/test_spacing.py @@ -27,7 +27,7 @@ from tests.lazy_transforms_utils import test_resampler_lazy from tests.test_utils import TEST_DEVICES, TEST_NDARRAYS_ALL, assert_allclose, dict_product, skip_if_quick -_TEMPLATE_5_COMPILED = torch.tensor([[[[0.75, 0.75]], [[0.75, 0.75]], [[0.75, 0.75]]]]) +_TEMPLATE_5_COMPILED = torch.tensor([[[[1.0, 1.0]], [[1.0, 1.0]], [[1.0, 1.0]]]]) _TEMPLATE_5_NATIVE = torch.tensor([[[[0.95527864, 0.95527864]], [[1.0, 1.0]], [[1.0, 1.0]]]])