Prevent implicit conversion of metatensor to numpy array

Signed-off-by: Davis Vigneault <[email protected]>

Author: sudomakeinstall

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