Merge branch 'dev' into 4980-get-wsi-at-mpp

Author: ericspod

monai/networks/layers/spatial_transforms.py

@@ -566,7 +566,7 @@ def forward(
                 affine=theta,
                 src_size=src_size[2:],
                 dst_size=dst_size[2:],
-                align_corners=False,
+                align_corners=self.align_corners,
                 zero_centered=self.zero_centered,
             )
         if self.reverse_indexing:

monai/transforms/compose.py

@@ -37,6 +37,31 @@
 __all__ = ["Compose", "OneOf", "RandomOrder", "SomeOf", "execute_compose"]
 
 
+def _inverse_one(
+    t: InvertibleTransform, data: Any, map_items: bool | int, unpack_items: bool, log_stats: bool | str
+) -> Any:
+    """Invert a single transform, delegating directly to nested ``Compose`` objects.
+
+    When ``t`` is a ``Compose`` instance its own ``inverse()`` is called so that
+    the child's ``map_items`` setting is respected.  For all other invertible
+    transforms, ``apply_transform`` is used with ``lazy=False``.
+
+    Args:
+        t: The invertible transform to invert.
+        data: Data to be inverted.
+        map_items: Whether to map over list/tuple items (forwarded to
+            ``apply_transform`` for non-``Compose`` transforms).
+        unpack_items: Whether to unpack data as parameters.
+        log_stats: Logger name or boolean for logging.
+
+    Returns:
+        The inverted data.
+    """
+    if isinstance(t, Compose):
+        return t.inverse(data)
+    return apply_transform(t.inverse, data, map_items, unpack_items, lazy=False, log_stats=log_stats)
+
+
 def execute_compose(
     data: NdarrayOrTensor | Sequence[NdarrayOrTensor] | Mapping[Any, NdarrayOrTensor],
     transforms: Sequence[Any],
@@ -315,20 +340,32 @@ def get_index_of_first(self, predicate):
         return None
 
     def flatten(self):
-        """Return a Composition with a simple list of transforms, as opposed to any nested Compositions.
+        """Return a Composition with a flattened list of transforms.
+
+        Nested ``Compose`` objects that share the same ``map_items`` setting as
+        the parent are inlined.  Nested ``Compose`` objects with a *different*
+        ``map_items`` value are kept as-is so their item-mapping behaviour is
+        preserved at runtime and during inversion.
 
         e.g., `t1 = Compose([x, x, x, x, Compose([Compose([x, x]), x, x])]).flatten()`
         will result in the equivalent of `t1 = Compose([x, x, x, x, x, x, x, x])`.
 
         """
         new_transforms = []
         for t in self.transforms:
-            if type(t) is Compose:  # nopep8
+            if type(t) is Compose and t.map_items == self.map_items:
                 new_transforms += t.flatten().transforms
             else:
                 new_transforms.append(t)
 
-        return Compose(new_transforms)
+        return Compose(
+            new_transforms,
+            map_items=self.map_items,
+            unpack_items=self.unpack_items,
+            log_stats=self.log_stats,
+            lazy=self._lazy,
+            overrides=self.overrides,
+        )
 
     def __len__(self):
         """Return number of transformations."""
@@ -365,9 +402,7 @@ def inverse(self, data):
             )
         # loop backwards over transforms
         for t in reversed(invertible_transforms):
-            data = apply_transform(
-                t.inverse, data, self.map_items, self.unpack_items, lazy=False, log_stats=self.log_stats
-            )
+            data = _inverse_one(t, data, self.map_items, self.unpack_items, self.log_stats)
         return data
 
     @staticmethod
@@ -622,9 +657,7 @@ def inverse(self, data):
         # loop backwards over transforms
         for o in reversed(applied_order):
             if isinstance(self.transforms[o], InvertibleTransform):
-                data = apply_transform(
-                    self.transforms[o].inverse, data, self.map_items, self.unpack_items, log_stats=self.log_stats
-                )
+                data = _inverse_one(self.transforms[o], data, self.map_items, self.unpack_items, self.log_stats)
         return data
 
 
@@ -789,8 +822,6 @@ def inverse(self, data):
         # loop backwards over transforms
         for o in reversed(applied_order):
             if isinstance(self.transforms[o], InvertibleTransform):
-                data = apply_transform(
-                    self.transforms[o].inverse, data, self.map_items, self.unpack_items, log_stats=self.log_stats
-                )
+                data = _inverse_one(self.transforms[o], data, self.map_items, self.unpack_items, self.log_stats)
 
         return data

monai/transforms/io/array.py

@@ -299,7 +299,7 @@ def __call__(self, filename: Sequence[PathLike] | PathLike, reader: ImageReader
             img_array, meta_data, self.simple_keys, pattern=self.pattern, sep=self.sep
         )
         if self.ensure_channel_first:
-            img = EnsureChannelFirst()(img)
+            img = EnsureChannelFirst()(img, meta_dict=meta_data)
         if self.image_only:
             return img
         return img, img.meta if isinstance(img, MetaTensor) else meta_data

monai/transforms/lazy/functional.py

@@ -16,6 +16,7 @@
 
 import torch
 
+import monai
 from monai.apps.utils import get_logger
 from monai.config import NdarrayOrTensor
 from monai.data.meta_tensor import MetaTensor
@@ -29,7 +30,7 @@
 )
 from monai.transforms.traits import LazyTrait
 from monai.transforms.transform import MapTransform
-from monai.utils import LazyAttr, look_up_option
+from monai.utils import LazyAttr, TraceKeys, look_up_option
 
 __all__ = ["apply_pending_transforms", "apply_pending_transforms_in_order", "apply_pending"]
 
@@ -289,6 +290,25 @@ def apply_pending(data: torch.Tensor | MetaTensor, pending: list | None = None,
         cumulative_xform = combine_transforms(cumulative_xform, next_matrix)
         cur_kwargs.update(new_kwargs)
     cur_kwargs.update(override_kwargs)
+    if len(pending) == 1 and isinstance(pending[0], dict):
+        p0 = pending[0]
+        extra_info = p0.get(TraceKeys.EXTRA_INFO)
+        align_corners = cur_kwargs.get(LazyAttr.ALIGN_CORNERS, False)
+        if (
+            isinstance(extra_info, dict)
+            and "affine" in extra_info
+            and TraceKeys.ORIG_SIZE in p0
+            and align_corners not in (False, TraceKeys.NONE)
+            and not isinstance(cur_kwargs.get(LazyAttr.INTERP_MODE), int)
+        ):
+            out_size = cur_kwargs.get(LazyAttr.SHAPE, p0.get(LazyAttr.SHAPE, p0[TraceKeys.ORIG_SIZE]))
+            cumulative_xform = monai.transforms.Affine.compute_w_affine(
+                len(tuple(p0[TraceKeys.ORIG_SIZE])),
+                extra_info["affine"],
+                p0[TraceKeys.ORIG_SIZE],
+                out_size,
+                align_corners=True,
+            )
     data = resample(data.to(device), cumulative_xform, cur_kwargs)
     if isinstance(data, MetaTensor):
         for p in pending:

monai/transforms/lazy/utils.py

@@ -20,7 +20,7 @@
 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, convert_to_numpy, convert_to_tensor, look_up_option
+from monai.utils import LazyAttr, TraceKeys, convert_to_numpy, convert_to_tensor, look_up_option
 
 __all__ = ["resample", "combine_transforms"]
 
@@ -90,7 +90,11 @@ def affine_from_pending(pending_item):
 
 
 def kwargs_from_pending(pending_item):
-    """Extract kwargs from a pending transform 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.
+    """
     if not isinstance(pending_item, dict):
         return {}
     ret = {
@@ -101,7 +105,13 @@ def kwargs_from_pending(pending_item):
         ret[LazyAttr.SHAPE] = pending_item[LazyAttr.SHAPE]
     if LazyAttr.DTYPE in pending_item:
         ret[LazyAttr.DTYPE] = pending_item[LazyAttr.DTYPE]
-    return ret  # adding support of pending_item['extra_info']??
+    # 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:
+        align_corners_val = extra_info["align_corners"]
+        if isinstance(align_corners_val, bool):
+            ret[LazyAttr.ALIGN_CORNERS] = align_corners_val
+    return ret
 
 
 def is_compatible_apply_kwargs(kwargs_1, kwargs_2):

monai/transforms/spatial/array.py

@@ -540,7 +540,8 @@ def __call__(
         if self.recompute_affine and isinstance(data_array, MetaTensor):
             if lazy_:
                 raise NotImplementedError("recompute_affine is not supported with lazy evaluation.")
-            a = scale_affine(original_spatial_shape, actual_shape)
+            ac = align_corners if align_corners is not None else self.sp_resample.align_corners
+            a = scale_affine(original_spatial_shape, actual_shape, align_corners=ac)
             data_array.affine = convert_to_dst_type(a, affine_)[0]  # type: ignore
         return data_array
 
@@ -2322,12 +2323,22 @@ def __call__(
         )
 
     @classmethod
-    def compute_w_affine(cls, spatial_rank, mat, img_size, sp_size):
+    def compute_w_affine(cls, spatial_rank, mat, img_size, sp_size, align_corners: bool = False):
         r = int(spatial_rank)
         mat = to_affine_nd(r, mat)
         shift_1 = create_translate(r, [float(d - 1) / 2 for d in img_size[:r]])
         shift_2 = create_translate(r, [-float(d - 1) / 2 for d in sp_size[:r]])
-        mat = shift_1 @ convert_data_type(mat, np.ndarray)[0] @ shift_2
+        mat = convert_data_type(mat, np.ndarray)[0]
+        if align_corners:
+            # Keep lazy world-affine consistent with eager sampling:
+            # x_in = T_in @ S_in^-1 @ A_centered @ S_out @ T_out^-1 @ x_out
+            src_scale = create_scale(r, [(max(float(d), 2.0) - 1.0) / max(float(d), 2.0) for d in img_size[:r]])
+            dst_scale = create_scale(r, [max(float(d), 2.0) / (max(float(d), 2.0) - 1.0) for d in sp_size[:r]])
+            src_scale = convert_data_type(src_scale, np.ndarray)[0]
+            dst_scale = convert_data_type(dst_scale, np.ndarray)[0]
+            mat = shift_1 @ src_scale @ mat @ dst_scale @ shift_2
+        else:
+            mat = shift_1 @ mat @ shift_2
         return mat
 
     def inverse(self, data: torch.Tensor) -> torch.Tensor:

monai/transforms/spatial/functional.py

@@ -304,7 +304,7 @@ def resize(
     meta_info = TraceableTransform.track_transform_meta(
         img,
         sp_size=out_size,
-        affine=scale_affine(orig_size, out_size),
+        affine=scale_affine(orig_size, out_size, align_corners=align_corners if align_corners is not None else False),
         extra_info=extra_info,
         orig_size=orig_size,
         transform_info=transform_info,
@@ -439,7 +439,7 @@ def zoom(img, scale_factor, keep_size, mode, padding_mode, align_corners, dtype,
     """
     im_shape = img.peek_pending_shape() if isinstance(img, MetaTensor) else img.shape[1:]
     output_size = [int(math.floor(float(i) * z)) for i, z in zip(im_shape, scale_factor)]
-    xform = scale_affine(im_shape, output_size)
+    xform = scale_affine(im_shape, output_size, align_corners=align_corners if align_corners is not None else False)
     extra_info = {
         "mode": mode,
         "align_corners": align_corners if align_corners is not None else TraceKeys.NONE,

monai/transforms/transform.py

@@ -143,10 +143,13 @@ def apply_transform(
     try:
         map_items_ = int(map_items) if isinstance(map_items, bool) else map_items
         if isinstance(data, (list, tuple)) and map_items_ > 0 and not isinstance(transform, ReduceTrait):
-            return [
-                apply_transform(transform, item, map_items_ - 1, unpack_items, log_stats, lazy, overrides)
-                for item in data
-            ]
+            # If the transform is a Compose with its own map_items, let it handle list/tuple
+            # expansion internally so that nested Compose map_items settings are respected.
+            if not isinstance(transform, transforms.compose.Compose):
+                return [
+                    apply_transform(transform, item, map_items_ - 1, unpack_items, log_stats, lazy, overrides)
+                    for item in data
+                ]
         return _apply_transform(transform, data, unpack_items, lazy, overrides, log_stats)
     except Exception as e:
         # if in debug mode, don't swallow exception so that the breakpoint

monai/transforms/utils.py

@@ -2097,14 +2097,16 @@ def convert_to_contiguous(
         return data
 
 
-def scale_affine(spatial_size, new_spatial_size, centered: bool = True):
+def scale_affine(spatial_size, new_spatial_size, centered: bool = True, align_corners: bool = False):
     """
     Compute the scaling matrix according to the new spatial size
 
     Args:
         spatial_size: original spatial size.
         new_spatial_size: new spatial size.
         centered: whether the scaling is with respect to the image center (True, default) or corner (False).
+            Ignored when ``align_corners=True``, since corner-aligned scaling is inherently centered.
+        align_corners: if True, use (size-1) based scaling to match torch.nn.functional.interpolate behavior.
 
     Returns:
         the scaling matrix.
@@ -2113,9 +2115,19 @@ def scale_affine(spatial_size, new_spatial_size, centered: bool = True):
     r = max(len(new_spatial_size), len(spatial_size))
     if spatial_size == new_spatial_size:
         return np.eye(r + 1)
-    s = np.array([float(o) / float(max(n, 1)) for o, n in zip(spatial_size, new_spatial_size)], dtype=float)
+    if align_corners:
+        # Match interpolate behavior: (src-1)/(dst-1); when dst == 1 the scale collapses to 0
+        s = np.array(
+            [0.0 if float(n) == 1 else (float(o) - 1) / (float(n) - 1) for o, n in zip(spatial_size, new_spatial_size)],
+            dtype=float,
+        )
+    else:
+        # Standard scaling: src/dst
+        s = np.array([float(o) / float(max(n, 1)) for o, n in zip(spatial_size, new_spatial_size)], dtype=float)
     scale = create_scale(r, s.tolist())
-    if centered:
+    if centered and not align_corners:
+        # For align_corners=False, add offset to center the scaling
+        # For align_corners=True, the scaling is inherently centered (corners map to corners)
         scale[:r, -1] = (np.diag(scale)[:r] - 1) / 2.0  # type: ignore
     return scale
 

tests/networks/layers/test_affine_transform.py

@@ -154,21 +154,21 @@ def test_zoom_1(self):
         affine = torch.as_tensor([[2.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
         image = torch.arange(1.0, 13.0).view(1, 1, 3, 4).to(device=torch.device("cpu:0"))
         out = AffineTransform()(image, affine, (1, 4))
-        expected = [[[[2.333333, 3.333333, 4.333333, 5.333333]]]]
+        expected = [[[[5.0, 6.0, 7.0, 8.0]]]]
         np.testing.assert_allclose(out, expected, atol=_rtol)
 
     def test_zoom_2(self):
         affine = torch.as_tensor([[2.0, 0.0, 0.0], [0.0, 2.0, 0.0]], dtype=torch.float32)
         image = torch.arange(1.0, 13.0).view(1, 1, 3, 4).to(device=torch.device("cpu:0"))
         out = AffineTransform((1, 2))(image, affine)
-        expected = [[[[1.458333, 4.958333]]]]
+        expected = [[[[5.0, 7.0]]]]
         np.testing.assert_allclose(out, expected, atol=1e-5, rtol=_rtol)
 
     def test_zoom_zero_center(self):
         affine = torch.as_tensor([[2.0, 0.0, 0.0], [0.0, 2.0, 0.0]], dtype=torch.float32)
         image = torch.arange(1.0, 13.0).view(1, 1, 3, 4).to(device=torch.device("cpu:0"))
         out = AffineTransform((1, 2), zero_centered=True)(image, affine)
-        expected = [[[[5.5, 7.5]]]]
+        expected = [[[[5.0, 8.0]]]]
         np.testing.assert_allclose(out, expected, atol=1e-5, rtol=_rtol)
 
     def test_affine_transform_minimum(self):
@@ -380,6 +380,55 @@ def test_forward_3d(self):
         np.testing.assert_allclose(actual, expected)
         np.testing.assert_allclose(list(theta.shape), [1, 3, 4])
 
+    def test_align_corners_consistency(self):
+        """
+        Test that align_corners is consistently used between to_norm_affine and grid_sample.
+
+        With an identity affine transform, the output should match the input regardless of
+        the align_corners setting. This test verifies that the coordinate normalization
+        in to_norm_affine uses the same align_corners value as affine_grid/grid_sample.
+        """
+        # Create a simple test image
+        image = torch.arange(1.0, 13.0).view(1, 1, 3, 4)
+
+        # Identity affine in pixel space (i, j, k convention with reverse_indexing=True)
+        identity_affine = torch.eye(3, dtype=torch.float32).unsqueeze(0)
+
+        # Test with align_corners=True (the default)
+        xform_true = AffineTransform(align_corners=True)
+        out_true = xform_true(image, identity_affine)
+        np.testing.assert_allclose(out_true.detach().cpu().numpy(), image.detach().cpu().numpy(), atol=1e-5, rtol=_rtol)
+
+        # Test with align_corners=False
+        xform_false = AffineTransform(align_corners=False)
+        out_false = xform_false(image, identity_affine)
+        np.testing.assert_allclose(
+            out_false.detach().cpu().numpy(), image.detach().cpu().numpy(), atol=1e-5, rtol=_rtol
+        )
+
+    def test_align_corners_true_translation(self):
+        """
+        Test that translation works correctly with align_corners=True.
+
+        This ensures to_norm_affine correctly converts pixel-space translations
+        to normalized coordinates when align_corners=True.
+        """
+        # 4x4 image
+        image = torch.arange(1.0, 17.0).view(1, 1, 4, 4)
+
+        # Translate by +1 pixel in the j direction (column direction)
+        # With reverse_indexing=True (default), this is the last spatial dimension
+        # Positive translation in the affine shifts the sampling grid, resulting in
+        # the output appearing shifted in the opposite direction
+        affine = torch.tensor([[[1.0, 0.0, 0.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]]])
+
+        xform = AffineTransform(align_corners=True, padding_mode="zeros")
+        out = xform(image, affine)
+
+        # Expected: shift columns left by 1, rightmost column becomes 0
+        expected = torch.tensor([[[[2, 3, 4, 0], [6, 7, 8, 0], [10, 11, 12, 0], [14, 15, 16, 0]]]], dtype=torch.float32)
+        np.testing.assert_allclose(out.detach().cpu().numpy(), expected.detach().cpu().numpy(), atol=1e-4, rtol=_rtol)
+
 
 if __name__ == "__main__":
     unittest.main()