From 2b1e3d0def8cf102ab6164a7a13da40ef525bf42 Mon Sep 17 00:00:00 2001
From: YunLiu <55491388+KumoLiu@users.noreply.github.com>
Date: Wed, 11 Oct 2023 20:03:45 +0800
Subject: [PATCH] Add support for metatensor in AddGuidanceFromPointsDeepEditd
 and ResizeGuidanceMultipleLabelDeepEditd (#7115)

Fixes #7114

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [ ] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

---------

Signed-off-by: KumoLiu <yunl@nvidia.com>
---
 monai/apps/deepedit/transforms.py | 30 +++++++++++++++++++++++++-----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/monai/apps/deepedit/transforms.py b/monai/apps/deepedit/transforms.py
index 880f90baae..14060c6b0f 100644
--- a/monai/apps/deepedit/transforms.py
+++ b/monai/apps/deepedit/transforms.py
@@ -662,13 +662,21 @@ def _apply(clicks, factor):
     def __call__(self, data):
         d = dict(data)
         meta_dict_key = self.meta_keys or f"{self.ref_image}_{self.meta_key_postfix}"
-        if meta_dict_key not in d:
-            raise RuntimeError(f"Missing meta_dict {meta_dict_key} in data!")
-        if "spatial_shape" not in d[meta_dict_key]:
+        # extract affine matrix from metadata
+        if isinstance(d[self.ref_image], MetaTensor):
+            meta_dict = d[self.ref_image].meta  # type: ignore
+        elif meta_dict_key in d:
+            meta_dict = d[meta_dict_key]
+        else:
+            raise ValueError(
+                f"{meta_dict_key} is not found. Please check whether it is the correct the image meta key."
+            )
+
+        if "spatial_shape" not in meta_dict:
             raise RuntimeError('Missing "spatial_shape" in meta_dict!')
 
         # Assume channel is first and depth is last CHWD
-        original_shape = d[meta_dict_key]["spatial_shape"]
+        original_shape = meta_dict["spatial_shape"]
         current_shape = list(d[self.ref_image].shape)[1:]
 
         # in here we assume the depth dimension is in the last dimension of "original_shape" and "current_shape"
@@ -698,7 +706,19 @@ def __call__(self, data):
         d = dict(data)
         # Assume channel is first and depth is last CHWD
         current_shape = d[self.ref_image].shape[1:]
-        original_shape = d["image_meta_dict"]["spatial_shape"]
+
+        meta_dict_key = "image_meta_dict"
+        # extract affine matrix from metadata
+        if isinstance(d[self.ref_image], MetaTensor):
+            meta_dict = d[self.ref_image].meta  # type: ignore
+        elif meta_dict_key in d:
+            meta_dict = d[meta_dict_key]
+        else:
+            raise ValueError(
+                f"{meta_dict_key} is not found. Please check whether it is the correct the image meta key."
+            )
+
+        original_shape = meta_dict["spatial_shape"]
 
         factor = np.divide(current_shape, original_shape)
         all_guidances = {}