diff --git a/monai/transforms/io/array.py b/monai/transforms/io/array.py index 6fc278211f..0c0f5fcf62 100644 --- a/monai/transforms/io/array.py +++ b/monai/transforms/io/array.py @@ -510,21 +510,22 @@ def __call__( kw = self.fname_formatter(meta_data, self) filename = self.folder_layout.filename(**kw) + data_kwargs = self.data_kwargs.copy() if meta_data: meta_spatial_shape = ensure_tuple(meta_data.get("spatial_shape", ())) if len(meta_spatial_shape) >= len(img.shape): - self.data_kwargs["channel_dim"] = None - elif is_no_channel(self.data_kwargs.get("channel_dim")): + data_kwargs["channel_dim"] = None + elif is_no_channel(data_kwargs.get("channel_dim")): warnings.warn( f"data shape {img.shape} (with spatial shape {meta_spatial_shape}) " - f"but SaveImage `channel_dim` is set to {self.data_kwargs.get('channel_dim')} no channel." + f"but SaveImage `channel_dim` is set to {data_kwargs.get('channel_dim')} no channel." ) err = [] for writer_cls in self.writers: try: writer_obj = writer_cls(**self.init_kwargs) - writer_obj.set_data_array(data_array=img, **self.data_kwargs) + writer_obj.set_data_array(data_array=img, **data_kwargs) writer_obj.set_metadata(meta_dict=meta_data, **self.meta_kwargs) writer_obj.write(filename, **self.write_kwargs) self.writer_obj = writer_obj diff --git a/tests/transforms/test_save_image.py b/tests/transforms/test_save_image.py index ed7061095d..afa5ff1918 100644 --- a/tests/transforms/test_save_image.py +++ b/tests/transforms/test_save_image.py @@ -15,6 +15,7 @@ import tempfile import unittest +import numpy as np import torch from parameterized import parameterized @@ -23,6 +24,7 @@ from monai.utils import optional_import _, has_itk = optional_import("itk", allow_namespace_pkg=True) +nib, has_nibabel = optional_import("nibabel") TEST_CASE_1 = [torch.randint(0, 255, (1, 2, 3, 4)), {"filename_or_obj": "testfile0.nii.gz"}, ".nii.gz", False] @@ -76,5 +78,43 @@ def test_saved_content_with_filename(self, test_data, output_ext, resample): self.assertTrue(os.path.exists(filename + output_ext)) +@unittest.skipUnless(has_nibabel, "nibabel not installed") +class TestSaveImageChannelDim(unittest.TestCase): + + @parameterized.expand( + [ + ("numpy_single", "numpy", 1, 0, True), + ("numpy_no_metadata", "numpy", 1, 0, False), + ("tensor", "tensor", 1, 0, True), + ("metatensor", "metatensor", 1, 0, True), + ("multiple_channels", "numpy", 2, 0, True), + ("channel_last", "numpy", 2, -1, True), + ] + ) + def test_reuse_after_channel_less(self, _, input_type, channels, channel_dim, with_metadata): + spatial_shape = (2, 3, 4) + first = np.arange(24, dtype=np.float32).reshape(spatial_shape) + second = np.arange(channels * 24, dtype=np.float32).reshape((channels, *spatial_shape)) + 100 + expected = second[0] if channels == 1 else np.moveaxis(second, 0, -1) + second = np.moveaxis(second, 0, channel_dim) + metadata = {"spatial_shape": spatial_shape} + if input_type == "tensor": + first, second = torch.as_tensor(first), torch.as_tensor(second) + elif input_type == "metatensor": + first, second = MetaTensor(first, meta=metadata), MetaTensor(second, meta=metadata) + + with tempfile.TemporaryDirectory() as tempdir: + saver = SaveImage(writer="NibabelWriter", resample=False, print_log=False) + saver.set_options(data_kwargs={"channel_dim": channel_dim}) + original_options = saver.data_kwargs.copy() + saver(first, meta_data=metadata, filename=os.path.join(tempdir, "first")) + np.testing.assert_array_equal( + nib.load(os.path.join(tempdir, "first.nii.gz")).get_fdata(), np.arange(24).reshape(spatial_shape) + ) + saver(second, meta_data=metadata if with_metadata else None, filename=os.path.join(tempdir, "second")) + np.testing.assert_array_equal(nib.load(os.path.join(tempdir, "second.nii.gz")).get_fdata(), expected) + self.assertEqual(saver.data_kwargs, original_options) + + if __name__ == "__main__": unittest.main() diff --git a/tests/transforms/test_save_imaged.py b/tests/transforms/test_save_imaged.py index d2095a7554..af9f079774 100644 --- a/tests/transforms/test_save_imaged.py +++ b/tests/transforms/test_save_imaged.py @@ -16,6 +16,7 @@ import unittest from pathlib import Path +import numpy as np import torch from parameterized import parameterized @@ -26,6 +27,7 @@ from monai.utils import optional_import _, has_itk = optional_import("itk", allow_namespace_pkg=True) +nib, has_nibabel = optional_import("nibabel") TEST_CASE_1 = [ {"img": MetaTensor(torch.randint(0, 255, (1, 2, 3, 4)), meta={"filename_or_obj": "testfile0.nii.gz"})}, @@ -126,5 +128,40 @@ def test_includes_metadata(self, test_data, output_ext, resample): self.assertTrue(os.path.exists(test_data["img"].meta["saved_to"])) +@unittest.skipUnless(has_nibabel, "nibabel not installed") +class TestSaveImagedChannelDim(unittest.TestCase): + + @parameterized.expand( + [ + ("image_first", ("image", "label"), False), + ("label_first", ("label", "image"), False), + ("metatensor", ("image", "label"), True), + ] + ) + def test_mixed_channel_dimensions(self, _, keys, use_metatensor): + image = np.arange(24, dtype=np.float32).reshape(2, 3, 4) + label = (image + 100)[None] + with tempfile.TemporaryDirectory() as tempdir: + saver = SaveImaged( + keys=keys, + output_dir=tempdir, + writer="NibabelWriter", + resample=False, + separate_folder=False, + print_log=False, + ) + data = {} + for key, array in (("image", image), ("label", label)): + metadata = {"filename_or_obj": key, "spatial_shape": (2, 3, 4)} + data[key] = MetaTensor(array, meta=metadata) if use_metatensor else array + if not use_metatensor: + data[f"{key}_meta_dict"] = metadata + for _ in range(2): + saver(data) + for key, expected in (("image", image), ("label", label[0])): + saved = nib.load(os.path.join(tempdir, f"{key}_trans.nii.gz")).get_fdata() + np.testing.assert_array_equal(saved, expected) + + if __name__ == "__main__": unittest.main()