diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py index 53f1009a76..fbca906191 100644 --- a/monai/transforms/__init__.py +++ b/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, diff --git a/monai/transforms/spatial/array.py b/monai/transforms/spatial/array.py index 2ed2ae42c7..0fb96e88a1 100644 --- a/monai/transforms/spatial/array.py +++ b/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`. diff --git a/monai/transforms/spatial/dictionary.py b/monai/transforms/spatial/dictionary.py index eaff3be35d..02c922d52e 100644 --- a/monai/transforms/spatial/dictionary.py +++ b/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 diff --git a/tests/min_tests.py b/tests/min_tests.py index 426650eb04..e0710a93ec 100644 --- a/tests/min_tests.py +++ b/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", diff --git a/tests/test_resample_to_match.py b/tests/test_resample_to_match.py new file mode 100644 index 0000000000..20ec5597b4 --- /dev/null +++ b/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() diff --git a/tests/test_resample_to_matchd.py b/tests/test_resample_to_matchd.py new file mode 100644 index 0000000000..a727f619fa --- /dev/null +++ b/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()