fixes needed for metatensor changes in monai

src/segmantic/data/transforms.py

@@ -5,6 +5,7 @@
 import numpy as np
 from monai.config import KeysCollection, PathLike
 from monai.data.folder_layout import FolderLayout
+from monai.data.meta_tensor import MetaTensor
 from monai.transforms.transform import MapTransform
 from monai.utils import ImageMetaKey as Key
 from monai.utils import convert_to_numpy
@@ -125,22 +126,23 @@ def __call__(self, data: Mapping[Hashable, Any]) -> None:
         label_key = self.label_key if self.label_key in d else self.image_key
 
         image = d[image_key]
-        image = np.squeeze(convert_to_numpy(image))
         label = d[label_key]
-        label = np.squeeze(convert_to_numpy(label))
 
+        meta_key = f"{image_key}_{self.meta_key_postfix}"
+        meta_data = image.meta if isinstance(image, MetaTensor) else d[meta_key]
+
+        if "affine" in meta_data:
+            affine = convert_to_numpy(meta_data["affine"])
+        else:
+            affine = np.eye(4, 4)
+
+        image = np.squeeze(convert_to_numpy(image))
+        label = np.squeeze(convert_to_numpy(label))
         if image.shape != label.shape:
             raise RuntimeError(
                 f"{self.__class__.__name__}: image and label have different number of elements or shape"
             )
 
-        meta_key = f"{image_key}_{self.meta_key_postfix}"
-        if "affine" in d[meta_key]:
-            affine = convert_to_numpy(d[meta_key]["affine"])
-        else:
-            affine = np.eye(4, 4)
-
-        meta_data = d[meta_key] if meta_key in d else {}
         subject = meta_data[Key.FILENAME_OR_OBJ] if meta_data else str(self._data_index)
         self._data_index += 1
         filename = self.folder_layout.filename(subject=f"{subject}")

src/segmantic/detect/transforms.py

@@ -9,14 +9,15 @@
 from monai.config import KeysCollection, PathLike
 from monai.config.type_definitions import NdarrayOrTensor
 from monai.data.folder_layout import FolderLayout
+from monai.data.meta_tensor import MetaTensor
 from monai.transforms import (
     GaussianSmooth,
     ScaleIntensity,
     generate_spatial_bounding_box,
 )
 from monai.transforms.transform import MapTransform
 from monai.utils import ImageMetaKey as Key
-from monai.utils import convert_to_numpy, convert_to_tensor
+from monai.utils import convert_to_dst_type, convert_to_numpy, convert_to_tensor
 from monai.utils.enums import PostFix
 
 __all__ = ["LoadVert", "SaveVert", "EmbedVert", "ExtractVertPosition", "VertHeatMap"]
@@ -136,15 +137,20 @@ def __init__(
 
     def __call__(self, data):
         d = dict(data)
+
+        ref_image = d[self.ref_key]
+
         ref_meta_key = f"{self.ref_key}_{self.meta_key_postfix}"
-        if "affine" in d[ref_meta_key]:
-            affine = convert_to_numpy(d[ref_meta_key]["affine"])
+        meta_data = (
+            ref_image.meta if isinstance(ref_image, MetaTensor) else d[ref_meta_key]
+        )
+        if "affine" in meta_data:
+            affine = convert_to_numpy(meta_data["affine"])
         else:
             affine = np.eye(4, 4)
         rot_inv = np.linalg.inv(affine[:3, :3])
         t = affine[:3, 3]
 
-        ref_image = d[self.ref_key]
         for k in self.keys:
             vertices = d[k]
             out = np.zeros_like(ref_image)
@@ -153,10 +159,13 @@ def __call__(self, data):
                 continuous_idx = np.dot(rot_inv, p - t)
                 idx = np.round(continuous_idx).astype(int)
                 out[idx[0], idx[1], idx[2]] = label
-            d[k] = out
+            d[k] = convert_to_dst_type(out, dst=ref_image)[0]
 
-            meta_key = f"{k}_{self.meta_key_postfix}"
-            d[meta_key].update({"affine": affine, "original_channel_dim": "no_channel"})
+            if not isinstance(ref_image, MetaTensor):
+                meta_key = f"{k}_{self.meta_key_postfix}"
+                d[meta_key].update(
+                    {"affine": affine, "original_channel_dim": "no_channel"}
+                )
 
         return d