[3823] ResampleToMatch

monai/transforms/__init__.py

@@ -314,6 +314,7 @@
     RandRotate90,
     RandZoom,
     Resample,
+    ResampleToMatch,
     Resize,
     Rotate,
     Rotate90,
@@ -361,6 +362,9 @@
     RandZoomd,
     RandZoomD,
     RandZoomDict,
+    ResampleToMatchd,
+    ResampleToMatchD,
+    ResampleToMatchDict,
     Resized,
     ResizeD,
     ResizeDict,

monai/transforms/spatial/array.py

@@ -13,7 +13,8 @@
 https://github.com/Project-MONAI/MONAI/wiki/MONAI_Design
 """
 import warnings
-from typing import Any, Callable, List, Optional, Sequence, Tuple, Union
+from copy import deepcopy
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import torch
@@ -58,6 +59,7 @@
 
 __all__ = [
     "SpatialResample",
+    "ResampleToMatch",
     "Spacing",
     "Orientation",
     "Flip",
@@ -267,6 +269,51 @@ def __call__(
         return output_data, dst_affine
 
 
+class ResampleToMatch(SpatialResample):
+    """Resample an image to match given meta data. The affine matrix will be aligned,
+    and the size of the output image will match."""
+
+    def __call__(  # type: ignore
+        self,
+        img: NdarrayOrTensor,
+        src_meta: Optional[Dict] = None,
+        dst_meta: Optional[Dict] = None,
+        mode: Union[GridSampleMode, str, None] = GridSampleMode.BILINEAR,
+        padding_mode: Union[GridSamplePadMode, str, None] = GridSamplePadMode.BORDER,
+        align_corners: Optional[bool] = False,
+        dtype: DtypeLike = None,
+    ):
+        if src_meta is None:
+            raise RuntimeError("`in_meta` is missing")
+        if dst_meta is None:
+            raise RuntimeError("`out_meta` is missing")
+        mode = mode or self.mode
+        padding_mode = padding_mode or self.padding_mode
+        align_corners = self.align_corners if align_corners is None else align_corners
+        dtype = dtype or self.dtype
+        src_affine = src_meta.get("affine")
+        dst_affine = dst_meta.get("affine")
+        ndim = len(img.shape[1:])
+        spatial_size = dst_meta.get("dim", [])[1 : ndim + 2]
+        img, updated_affine = super().__call__(
+            img=img,
+            src_affine=src_affine,
+            dst_affine=dst_affine,
+            spatial_size=spatial_size,
+            mode=mode,
+            padding_mode=padding_mode,
+            align_corners=align_corners,
+            dtype=dtype,
+        )
+        dst_meta = deepcopy(dst_meta)
+        dst_meta["affine"] = updated_affine
+        if "dim" in dst_meta:
+            dst_meta["dim"] = src_meta.get("dim", [])
+        if "pixdim" in dst_meta:
+            dst_meta["pixdim"] = src_meta.get("pixdim", [])
+        return img, dst_meta
+
+
 class Spacing(Transform):
     """
     Resample input image into the specified `pixdim`.

monai/transforms/spatial/dictionary.py

@@ -43,6 +43,7 @@
     RandGridDistortion,
     RandRotate,
     RandZoom,
+    ResampleToMatch,
     Resize,
     Rotate,
     Rotate90,
@@ -71,6 +72,7 @@
 
 __all__ = [
     "SpatialResampled",
+    "ResampleToMatchd",
     "Spacingd",
     "Orientationd",
     "Rotate90d",
@@ -290,6 +292,72 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd
         return d
 
 
+class ResampleToMatchd(MapTransform, InvertibleTransform):
+    """Dictionary-based wrapper of :py:class:`monai.transforms.ResampleToMatch`."""
+
+    backend = ResampleToMatch.backend
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        template_key: str,
+        mode: GridSampleModeSequence = GridSampleMode.BILINEAR,
+        padding_mode: GridSamplePadModeSequence = GridSamplePadMode.BORDER,
+        align_corners: Union[Sequence[bool], bool] = False,
+        dtype: Union[Sequence[DtypeLike], DtypeLike] = np.float64,
+        allow_missing_keys: bool = False,
+    ):
+        """
+        Args:
+            keys: keys of the corresponding items to be transformed.
+            template_key: key to meta data that output should be resampled to match.
+            mode: {``"bilinear"``, ``"nearest"``}
+                Interpolation mode to calculate output values. Defaults to ``"bilinear"``.
+                See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
+                It also can be a sequence of string, each element corresponds to a key in ``keys``.
+            padding_mode: {``"zeros"``, ``"border"``, ``"reflection"``}
+                Padding mode for outside grid values. Defaults to ``"border"``.
+                See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
+                It also can be a sequence of string, each element corresponds to a key in ``keys``.
+            align_corners: Geometrically, we consider the pixels of the input as squares rather than points.
+                See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
+                It also can be a sequence of bool, each element corresponds to a key in ``keys``.
+            dtype: data type for resampling computation. Defaults to ``np.float64`` for best precision.
+                If None, use the data type of input data. To be compatible with other modules,
+                the output data type is always ``np.float32``.
+                It also can be a sequence of dtypes, each element corresponds to a key in ``keys``.
+            allow_missing_keys: don't raise exception if key is missing.
+        """
+        super().__init__(keys, allow_missing_keys)
+        self.template_key = template_key
+        self.mode = ensure_tuple_rep(mode, len(self.keys))
+        self.padding_mode = ensure_tuple_rep(padding_mode, len(self.keys))
+        self.align_corners = ensure_tuple_rep(align_corners, len(self.keys))
+        self.dtype = ensure_tuple_rep(dtype, len(self.keys))
+        self.resampler = ResampleToMatch()
+
+    def __call__(self, data):
+        d = deepcopy(dict(data))
+        dst_meta = d[self.template_key]
+        for (key, mode, padding_mode, align_corners, dtype) in self.key_iterator(
+            d, self.mode, self.padding_mode, self.align_corners, self.dtype
+        ):
+            src_meta_key = PostFix.meta(key)
+            src_meta = d[src_meta_key]
+            img, new_meta = self.resampler(
+                img=d[key],
+                src_meta=src_meta,
+                dst_meta=dst_meta,
+                mode=mode,
+                padding_mode=padding_mode,
+                align_corners=align_corners,
+                dtype=dtype,
+            )
+            d[key] = img
+            d[src_meta_key] = new_meta
+        return d
+
+
 class Spacingd(MapTransform, InvertibleTransform):
     """
     Dictionary-based wrapper of :py:class:`monai.transforms.Spacing`.
@@ -2035,6 +2103,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
 
 
 SpatialResampleD = SpatialResampleDict = SpatialResampled
+ResampleToMatchD = ResampleToMatchDict = ResampleToMatchd
 SpacingD = SpacingDict = Spacingd
 OrientationD = OrientationDict = Orientationd
 Rotate90D = Rotate90Dict = Rotate90d

tests/min_tests.py

@@ -131,6 +131,8 @@ def run_testsuit():
         "test_randtorchvisiond",
         "test_resize",
         "test_resized",
+        "test_resample_to_match",
+        "test_resample_to_matchd",
         "test_rotate",
         "test_rotated",
         "test_save_image",

tests/test_resample_to_match.py

@@ -0,0 +1,51 @@
+# 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.
+
+import os
+import tempfile
+import unittest
+
+from monai.transforms import Compose, EnsureChannelFirstd, LoadImaged, ResampleToMatch, SaveImaged
+from tests.utils import assert_allclose, download_url_or_skip_test
+
+
+class TestResampleToMatch(unittest.TestCase):
+    def test_correct(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            url_1 = (
+                "https://github.com/rcuocolo/PROSTATEx_masks/raw/master/Files/"
+                + "lesions/Images/T2/ProstateX-0000_t2_tse_tra_4.nii.gz"
+            )
+            url_2 = (
+                "https://github.com/rcuocolo/PROSTATEx_masks/raw/master/Files/"
+                + "lesions/Images/ADC/ProstateX-0000_ep2d_diff_tra_7.nii.gz"
+            )
+            fname_1 = os.path.join(temp_dir, "file1.nii.gz")
+            fname_2 = os.path.join(temp_dir, "file2.nii.gz")
+            md5_1 = "adb3f1c4db66a6481c3e4a2a3033c7d5"
+            md5_2 = "f12a11ad0ebb0b1876e9e010564745d2"
+            download_url_or_skip_test(url=url_1, filepath=fname_1, hash_val=md5_1)
+            download_url_or_skip_test(url=url_2, filepath=fname_2, hash_val=md5_2)
+
+            loader = Compose([LoadImaged(("im1", "im2")), EnsureChannelFirstd(("im1", "im2"))])
+            data = loader({"im1": fname_1, "im2": fname_2})
+
+            im_mod, meta = ResampleToMatch()(data["im2"], data["im2_meta_dict"], data["im1_meta_dict"])
+            # for visual inspection
+            saver = SaveImaged("im3", output_dir=temp_dir, output_postfix="", separate_folder=False)
+            meta["filename_or_obj"] = "file3.nii.gz"
+            saver({"im3": im_mod, "im3_meta_dict": meta})
+
+            assert_allclose(im_mod.shape, data["im1"].shape)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_resample_to_matchd.py

@@ -0,0 +1,58 @@
+# 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.
+
+import os
+import tempfile
+import unittest
+
+from monai.transforms import Compose, CopyItemsd, EnsureChannelFirstd, Lambda, LoadImaged, ResampleToMatchd, SaveImaged
+from tests.utils import assert_allclose, download_url_or_skip_test
+
+
+def update_fname(d):
+    d["im3_meta_dict"]["filename_or_obj"] = "file3.nii.gz"
+    return d
+
+
+class TestResampleToMatchd(unittest.TestCase):
+    def test_correct(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            url_1 = (
+                "https://github.com/rcuocolo/PROSTATEx_masks/raw/master/Files/"
+                + "lesions/Images/T2/ProstateX-0000_t2_tse_tra_4.nii.gz"
+            )
+            url_2 = (
+                "https://github.com/rcuocolo/PROSTATEx_masks/raw/master/Files/"
+                + "lesions/Images/ADC/ProstateX-0000_ep2d_diff_tra_7.nii.gz"
+            )
+            fname_1 = os.path.join(temp_dir, "file1.nii.gz")
+            fname_2 = os.path.join(temp_dir, "file2.nii.gz")
+            md5_1 = "adb3f1c4db66a6481c3e4a2a3033c7d5"
+            md5_2 = "f12a11ad0ebb0b1876e9e010564745d2"
+            download_url_or_skip_test(url=url_1, filepath=fname_1, hash_val=md5_1)
+            download_url_or_skip_test(url=url_2, filepath=fname_2, hash_val=md5_2)
+
+            transforms = Compose(
+                [
+                    LoadImaged(("im1", "im2")),
+                    EnsureChannelFirstd(("im1", "im2")),
+                    CopyItemsd(("im2", "im2_meta_dict"), names=("im3", "im3_meta_dict")),
+                    ResampleToMatchd("im3", "im1_meta_dict"),
+                    Lambda(update_fname),
+                    SaveImaged("im3", output_dir=temp_dir, output_postfix="", separate_folder=False),
+                ]
+            )
+            data = transforms({"im1": fname_1, "im2": fname_2})
+            assert_allclose(data["im1"].shape, data["im3"].shape)
+
+
+if __name__ == "__main__":
+    unittest.main()