Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion monai/data/image_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -324,7 +324,7 @@ def convert_to_channel_last(
data = data[..., 0, :]
# if desired, remove trailing singleton dimensions
while squeeze_end_dims and data.shape[-1] == 1:
data = np.squeeze(data, -1)
data = data.squeeze(-1)
if contiguous:
data = ascontiguousarray(data)
return data
Expand Down
11 changes: 9 additions & 2 deletions tests/data/test_itk_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
import numpy as np
import torch

from monai.data import ITKWriter
from monai.data import ITKWriter, MetaTensor
from monai.utils import optional_import

itk, has_itk = optional_import("itk")
Expand All @@ -45,7 +45,7 @@ def test_rgb(self):
with tempfile.TemporaryDirectory() as tempdir:
fname = os.path.join(tempdir, "testing.png")
writer = ITKWriter(output_dtype=np.uint8)
writer.set_data_array(np.arange(48).reshape(3, 4, 4), channel_dim=0)
writer.set_data_array(torch.arange(48).reshape(3, 4, 4), channel_dim=0)
writer.set_metadata({"spatial_shape": (5, 5)})
writer.write(fname)

Expand All @@ -64,6 +64,13 @@ def test_no_channel(self):
np.testing.assert_allclose(output.shape, (4, 4, 3))
np.testing.assert_allclose(output[1, 1], (5, 21, 37))

def test_metatensor_preserved(self):
data = MetaTensor(np.arange(48).reshape(3, 4, 4, 1), meta={"test_key": "test_value"})
writer = ITKWriter()
writer.set_data_array(data, channel_dim=-1, squeeze_end_dims=True)
self.assertIsInstance(writer.data_obj, MetaTensor)
self.assertEqual(writer.data_obj.meta.get("test_key"), "test_value")
Comment thread
sudomakeinstall marked this conversation as resolved.


if __name__ == "__main__":
unittest.main()
Loading