From af8ce517107a21fc13fcb33b4cdf9acfa0976f7d Mon Sep 17 00:00:00 2001 From: Thomas Kierski <54414492+ThomasKierski@users.noreply.github.com> Date: Mon, 12 Oct 2026 11:38:26 -0400 Subject: [PATCH] Fix lazy resampling not respecting transform interpolation mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description With `Compose(..., lazy=True)`, the interpolation mode set on each transform was ignored when pending operations were applied. A binary label processed with `mode="nearest"` (for example `Spacingd(keys="label", pixdim=1.3, mode="nearest")`) came out with soft, interpolated values. The same pipeline with `lazy=False` correctly returned `{0, 1}`. This was caused by: 1. `kwargs_from_pending` (`monai/transforms/lazy/utils.py`) only reads `lazy_interpolation_mode` and `lazy_padding_mode` from the top level of a pending operation. No transform writes those keys. Spatial transforms record their modes in `extra_info["mode"]` and `extra_info["padding_mode"]`. So the mode was always `None`, and `SpatialResample` fell back to its default (bilinear). 2. `apply_pending` (`monai/transforms/lazy/functional.py`) merges each pending op's kwargs with `cur_kwargs.update(new_kwargs)`. A later op that has no mode (e.g. `SpatialPadd`, `CenterSpatialCropd`) overwrote an earlier mode with `None`. Changes: - `kwargs_from_pending` now falls back to `extra_info["mode"]` / `extra_info["padding_mode"]` when the explicit lazy attributes are not set. Explicit lazy attributes still take precedence. Recorded modes are converted to values the lazy resampler supports: - grid-sample modes (`nearest`, `bilinear`, `bicubic`) and spline orders 0–5 are kept; - `nearest-exact` becomes `nearest`; `linear` and `trilinear` become `bilinear`; - modes with no grid-sample equivalent (e.g. `area`) and pad modes stored by pad transforms (e.g. `"constant"`) are ignored, which keeps the previous default behaviour for them. - `apply_pending` no longer lets `None` values from later pending ops overwrite kwargs set by earlier ones. - `Compose(overrides=...)` behaviour is unchanged and still takes priority. - The existing behaviour of `is_compatible_apply_kwargs` (always `True`) is unchanged. If a single key mixes modes across transforms, the last non-`None` mode still wins. ### Types of changes - [x] Non-breaking change (fix or new feature that would not break existing functionality). - [ ] Breaking change (fix or new feature that would cause existing functionality to change). - [x] New tests added to cover the changes. - [ ] Integration tests passed locally by running `./runtests.sh -f -u --net --coverage`. - [ ] Quick tests passed locally by running `./runtests.sh --quick --unittests --disttests`. - [x] In-line docstrings updated. - [ ] Documentation updated, tested `make html` command in the `docs/` folder. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Thomas Kierski <54414492+ThomasKierski@users.noreply.github.com> --- monai/transforms/lazy/functional.py | 3 +- monai/transforms/lazy/utils.py | 97 +++++++++++++-- .../compose/test_lazy_interp_mode.py | 113 ++++++++++++++++++ 3 files changed, 205 insertions(+), 8 deletions(-) create mode 100644 tests/transforms/compose/test_lazy_interp_mode.py diff --git a/monai/transforms/lazy/functional.py b/monai/transforms/lazy/functional.py index 4120dece388..650198dceab 100644 --- a/monai/transforms/lazy/functional.py +++ b/monai/transforms/lazy/functional.py @@ -290,7 +290,8 @@ def apply_pending(data: torch.Tensor | MetaTensor, pending: list | None = None, next_matrix = to_affine_nd(_rank, next_matrix) cumulative_xform = combine_transforms(cumulative_xform, next_matrix) - cur_kwargs.update(new_kwargs) + # don't let pending items without a mode (e.g. crop/pad) discard a mode set by an earlier transform + cur_kwargs.update({k: v for k, v in new_kwargs.items() if v is not None}) cur_kwargs.update(override_kwargs) if len(pending) == 1 and isinstance(pending[0], dict): p0 = pending[0] diff --git a/monai/transforms/lazy/utils.py b/monai/transforms/lazy/utils.py index 1f8506031a9..df8e0908840 100644 --- a/monai/transforms/lazy/utils.py +++ b/monai/transforms/lazy/utils.py @@ -12,6 +12,7 @@ from __future__ import annotations import warnings +from enum import Enum import numpy as np import torch @@ -20,7 +21,17 @@ from monai.config import NdarrayOrTensor from monai.data.utils import AFFINE_TOL from monai.transforms.utils_pytorch_numpy_unification import allclose -from monai.utils import LazyAttr, TraceKeys, convert_to_numpy, convert_to_tensor, look_up_option +from monai.utils import ( + GridSampleMode, + GridSamplePadMode, + InterpolateMode, + LazyAttr, + NdimageMode, + TraceKeys, + convert_to_numpy, + convert_to_tensor, + look_up_option, +) __all__ = ["resample", "combine_transforms"] @@ -89,25 +100,97 @@ def affine_from_pending(pending_item): return pending_item +def _lazy_interp_mode(mode): + """ + Convert an interpolation mode recorded by a spatial transform into one usable by the lazy resampler + (:py:class:`monai.transforms.SpatialResample`). + + ``"nearest-exact"`` is mapped to grid-sample ``"nearest"``, which samples at the same (half-pixel aligned) + source locations; results can differ only for source locations exactly halfway between two voxels, where + grid-sample rounds half to even and ``torch.nn.functional.interpolate`` rounds half up. + + Args: + mode: the interpolation mode recorded in a pending operation's ``extra_info``, e.g. a + :py:class:`monai.utils.GridSampleMode` or :py:class:`monai.utils.InterpolateMode` value, its string + form, or a spline interpolation order (``0``-``5``). + + Returns: + A grid-sample mode string or a spline order supported by the lazy resampler, or ``None`` if ``mode`` is + unset or has no equivalent (e.g. ``"area"``), in which case the resampler's default mode is used. + """ + if isinstance(mode, Enum): + mode = mode.value + if isinstance(mode, (int, np.integer)) and not isinstance(mode, bool): + return int(mode) if int(mode) in range(6) else None # spline order + if not isinstance(mode, str): + return None + if mode in {m.value for m in GridSampleMode}: + return mode + if mode == InterpolateMode.NEAREST_EXACT: + return GridSampleMode.NEAREST.value + if mode.endswith("linear"): # linear, bilinear, trilinear + return GridSampleMode.BILINEAR.value + return None + + +def _lazy_padding_mode(padding_mode): + """ + Convert a padding mode recorded by a spatial transform into one usable by the lazy resampler + (:py:class:`monai.transforms.SpatialResample`). + + Args: + padding_mode: the padding mode recorded in a pending operation's ``extra_info``, e.g. a + :py:class:`monai.utils.GridSamplePadMode` or :py:class:`monai.utils.NdimageMode` value or its string form. + + Returns: + ``padding_mode`` as a string if it is a supported grid-sample or ndimage padding mode, otherwise ``None``, + in which case the resampler's default padding mode is used. + """ + if isinstance(padding_mode, Enum): + padding_mode = padding_mode.value + valid_modes = {*(m.value for m in GridSamplePadMode), *(m.value for m in NdimageMode)} + if isinstance(padding_mode, str) and padding_mode in valid_modes: + return padding_mode + return None + + def kwargs_from_pending(pending_item): """Extract kwargs from a pending transform item. When ``pending_item`` is a dict, ``align_corners`` is also extracted from its ``extra_info`` entry (if present and boolean) so the lazy pipeline preserves the original transform's alignment. + Likewise, if the interpolation/padding modes are not set explicitly on the pending item, the ``mode`` + and ``padding_mode`` recorded in ``extra_info`` by the spatial transform (e.g. ``Spacing(mode="nearest")``) + are used, so that lazy resampling honours the mode requested by the user. + + Args: + pending_item: a pending operation, typically a dict pushed by a lazy transform (see + :py:meth:`monai.data.MetaTensor.push_pending_operation`). Non-dict items (e.g. a bare affine + matrix) carry no kwargs. + + Returns: + A dict of resampling kwargs keyed by :py:class:`monai.utils.LazyAttr` values. For a dict + ``pending_item`` it always contains ``LazyAttr.INTERP_MODE`` and ``LazyAttr.PADDING_MODE`` (``None`` if + unset or unsupported), plus ``LazyAttr.SHAPE``, ``LazyAttr.DTYPE`` and ``LazyAttr.ALIGN_CORNERS`` when + available. An empty dict is returned for non-dict items. """ if not isinstance(pending_item, dict): return {} - ret = { - LazyAttr.INTERP_MODE: pending_item.get(LazyAttr.INTERP_MODE, None), # interpolation mode - LazyAttr.PADDING_MODE: pending_item.get(LazyAttr.PADDING_MODE, None), # padding mode - } + extra_info = pending_item.get(TraceKeys.EXTRA_INFO) + extra_info = extra_info if isinstance(extra_info, dict) else {} + interp_mode = pending_item.get(LazyAttr.INTERP_MODE, None) + if interp_mode is None: + interp_mode = _lazy_interp_mode(extra_info.get("mode")) + padding_mode = pending_item.get(LazyAttr.PADDING_MODE, None) + if padding_mode is None: + padding_mode = _lazy_padding_mode(extra_info.get("padding_mode")) + ret = {LazyAttr.INTERP_MODE: interp_mode, LazyAttr.PADDING_MODE: padding_mode} if LazyAttr.SHAPE in pending_item: ret[LazyAttr.SHAPE] = pending_item[LazyAttr.SHAPE] if LazyAttr.DTYPE in pending_item: ret[LazyAttr.DTYPE] = pending_item[LazyAttr.DTYPE] # Extract align_corners from extra_info if available - extra_info = pending_item.get(TraceKeys.EXTRA_INFO) - if isinstance(extra_info, dict) and "align_corners" in extra_info: + if "align_corners" in extra_info: align_corners_val = extra_info["align_corners"] if isinstance(align_corners_val, bool): ret[LazyAttr.ALIGN_CORNERS] = align_corners_val diff --git a/tests/transforms/compose/test_lazy_interp_mode.py b/tests/transforms/compose/test_lazy_interp_mode.py new file mode 100644 index 00000000000..0c1542ebee2 --- /dev/null +++ b/tests/transforms/compose/test_lazy_interp_mode.py @@ -0,0 +1,113 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import unittest + +import torch +from parameterized import parameterized + +from monai.data import MetaTensor +from monai.transforms import ( + CenterSpatialCropd, + Compose, + RandAffined, + Resize, + Resized, + ResizeWithPadOrCropd, + Rotated, + Spacingd, + SpatialPadd, + Zoomd, +) +from monai.transforms.lazy.utils import kwargs_from_pending +from monai.utils import LazyAttr, TraceKeys +from tests.test_utils import assert_allclose + +TEST_CASES = [ + ["spacing", [Spacingd("label", pixdim=1.3, mode="nearest")]], + ["spacing_pad", [Spacingd("label", pixdim=1.3, mode="nearest"), SpatialPadd("label", 40)]], + ["pad_spacing", [SpatialPadd("label", 40), Spacingd("label", pixdim=1.3, mode="nearest")]], + ["spacing_crop", [Spacingd("label", pixdim=0.7, mode="nearest"), CenterSpatialCropd("label", 30)]], + ["rotate_padcrop", [Rotated("label", 0.3, mode="nearest"), ResizeWithPadOrCropd("label", 36)]], + [ + "zoom_affine", + [Zoomd("label", 1.2, mode="nearest"), RandAffined("label", prob=1.0, rotate_range=0.3, mode="nearest")], + ], + ["resize", [Resized("label", (20, 20, 20), mode="nearest")]], +] + + +class TestLazyInterpMode(unittest.TestCase): + """Tests that lazy resampling honours the interpolation/padding modes set on each transform.""" + + @parameterized.expand(TEST_CASES) + def test_binary_label_stays_binary(self, _, transforms): + """Check a binary label resampled lazily with ``mode="nearest"`` contains only 0 and 1. + + Args: + _: name of the test case. + transforms: dictionary transforms (all using ``mode="nearest"`` where applicable) to compose lazily. + """ + torch.manual_seed(0) + label = MetaTensor((torch.rand(1, 32, 32, 32) > 0.5).float(), affine=torch.eye(4)) + pipeline = Compose(transforms, lazy=True) + pipeline.set_random_state(0) + out = pipeline({"label": label})["label"] + self.assertTrue(set(torch.unique(out).tolist()) <= {0.0, 1.0}) + + def test_bilinear_key_still_interpolates(self): + """Check per-key modes in one lazy pipeline: the bilinear key is interpolated, the nearest key stays binary.""" + torch.manual_seed(0) + data = { + "img": MetaTensor((torch.rand(1, 32, 32, 32) > 0.5).float(), affine=torch.eye(4)), + "label": MetaTensor((torch.rand(1, 32, 32, 32) > 0.5).float(), affine=torch.eye(4)), + } + keys = ["img", "label"] + pipeline = Compose([Spacingd(keys, pixdim=1.3, mode=("bilinear", "nearest")), SpatialPadd(keys, 40)], lazy=True) + out = pipeline(data) + self.assertTrue(set(torch.unique(out["label"]).tolist()) <= {0.0, 1.0}) + # a binary input only yields values strictly between 0 and 1 if it was actually interpolated + self.assertTrue(((out["img"] > 0) & (out["img"] < 1)).any()) + + def test_nearest_exact_resize_matches_eager(self): + """Check lazy ``Resize(mode="nearest-exact")`` matches eager output when no sample lies on an exact tie.""" + torch.manual_seed(0) + img = MetaTensor(torch.randint(0, 5, (1, 17, 23, 29)).float(), affine=torch.eye(4)) + eager = Resize((20, 20, 20), mode="nearest-exact")(img.clone()) + lazy = Compose([Resize((20, 20, 20), mode="nearest-exact")], lazy=True)(img.clone()) + assert_allclose(lazy, eager, type_test=False) + + def test_kwargs_from_pending_extra_info(self): + """Check ``kwargs_from_pending`` reads, prioritises and normalises modes recorded in ``extra_info``.""" + kw = kwargs_from_pending({TraceKeys.EXTRA_INFO: {"mode": "nearest", "padding_mode": "border"}}) + self.assertEqual(kw[LazyAttr.INTERP_MODE], "nearest") + self.assertEqual(kw[LazyAttr.PADDING_MODE], "border") + # explicit lazy attributes take precedence + kw = kwargs_from_pending({LazyAttr.INTERP_MODE: "bilinear", TraceKeys.EXTRA_INFO: {"mode": "nearest"}}) + self.assertEqual(kw[LazyAttr.INTERP_MODE], "bilinear") + # padding modes recorded by pad transforms are not interpolation modes + kw = kwargs_from_pending({TraceKeys.EXTRA_INFO: {"padded": ((0, 0),), "mode": "constant"}}) + self.assertIsNone(kw[LazyAttr.INTERP_MODE]) + # torch interpolate-only modes are mapped/ignored + self.assertEqual( + kwargs_from_pending({TraceKeys.EXTRA_INFO: {"mode": "nearest-exact"}})[LazyAttr.INTERP_MODE], "nearest" + ) + self.assertEqual( + kwargs_from_pending({TraceKeys.EXTRA_INFO: {"mode": "trilinear"}})[LazyAttr.INTERP_MODE], "bilinear" + ) + self.assertIsNone(kwargs_from_pending({TraceKeys.EXTRA_INFO: {"mode": "area"}})[LazyAttr.INTERP_MODE]) + self.assertEqual(kwargs_from_pending({TraceKeys.EXTRA_INFO: {"mode": 0}})[LazyAttr.INTERP_MODE], 0) + + +if __name__ == "__main__": + unittest.main()