From b0ca3921de98612b30c1d5c1e883637a7b072b1b Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Wed, 15 Mar 2023 16:20:18 +0000 Subject: [PATCH] update channel_dim None case Signed-off-by: Wenqi Li --- monai/data/image_writer.py | 9 ++++++--- tests/test_itk_writer.py | 11 +++++++++++ 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/monai/data/image_writer.py b/monai/data/image_writer.py index b467a5ae99..94b31c164c 100644 --- a/monai/data/image_writer.py +++ b/monai/data/image_writer.py @@ -414,9 +414,12 @@ def set_data_array( spatial_ndim=kwargs.pop("spatial_ndim", 3), contiguous=kwargs.pop("contiguous", True), ) - self.channel_dim = ( - channel_dim if self.data_obj is not None and len(self.data_obj.shape) >= _r else None - ) # channel dim is at the end + if channel_dim is None: + self.channel_dim = -1 # after convert_to_channel_last the last dim is the channel + else: + self.channel_dim = ( + channel_dim if self.data_obj is not None and len(self.data_obj.shape) >= _r else None + ) # channel dim is at the end def set_metadata(self, meta_dict: Mapping | None = None, resample: bool = True, **options): """ diff --git a/tests/test_itk_writer.py b/tests/test_itk_writer.py index 869ec7b947..75f848ec01 100644 --- a/tests/test_itk_writer.py +++ b/tests/test_itk_writer.py @@ -52,6 +52,17 @@ def test_rgb(self): np.testing.assert_allclose(output.shape, (5, 5, 3)) np.testing.assert_allclose(output[1, 1], (5, 5, 4)) + def test_no_channel(self): + with tempfile.TemporaryDirectory() as tempdir: + fname = os.path.join(tempdir, "testing.nii.gz") + writer = ITKWriter(output_dtype=np.uint8) + writer.set_data_array(np.arange(48).reshape(3, 4, 4), channel_dim=None) + writer.write(fname) + + output = np.asarray(itk.imread(fname)) + np.testing.assert_allclose(output.shape, (4, 3, 4)) + np.testing.assert_allclose(output[1, 1], (20, 21, 22, 23)) + if __name__ == "__main__": unittest.main()