Merge branch 'dev' into 8627-perceptual-loss-errors-out-after-hitting-the-maximum-number-of-downloads

Author: ericspod

monai/data/image_writer.py

@@ -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

tests/data/test_itk_writer.py

@@ -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")
@@ -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)
 
@@ -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")
+
 
 if __name__ == "__main__":
     unittest.main()