Fix lazy resampling not respecting transform interpolation mode

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]

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

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()