Keep SaveImage channel options local to each call

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

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

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