update deepedit

Signed-off-by: Yiheng Wang <vennw@nvidia.com>

ci/unit_tests/test_spleen_deepedit_annotation.py

@@ -11,6 +11,7 @@
 
 import os
 import shutil
+import sys
 import tempfile
 import unittest
 
@@ -109,6 +110,7 @@ def test_eval_config(self, override):
     @parameterized.expand([TEST_CASE_2])
     def test_infer_config(self, override):
         override["dataset_dir"] = self.dataset_dir
+        override["use_click"] = False
         bundle_root = override["bundle_root"]
 
         inferrer = ConfigWorkflow(
@@ -128,7 +130,7 @@ def test_infer_click_config(self, override):
             "dataset#data"
         ] = "$[{'image': i, 'background': [], 'spleen': [[6, 6, 6], [8, 8, 8]]} for i in @datalist]"
         bundle_root = override["bundle_root"]
-        print(override)
+        sys.path = [bundle_root] + sys.path
 
         inferrer = ConfigWorkflow(
             workflow="infer",

models/spleen_deepedit_annotation/configs/inference.json

@@ -19,7 +19,7 @@
     ],
     "number_intensity_ch": 1,
     "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')",
-    "use_click": false,
+    "use_click": true,
     "network_def": {
         "_target_": "DynUNet",
         "spatial_dims": 3,
@@ -101,6 +101,11 @@
         }
     ],
     "deepedit_transforms": [
+        {
+            "_target_": "scripts.transforms.OrientationGuidanceMultipleLabelDeepEditd",
+            "ref_image": "image",
+            "label_names": "@label_names"
+        },
         {
             "_target_": "AddGuidanceFromPointsDeepEditd",
             "ref_image": "image",

models/spleen_deepedit_annotation/configs/metadata.json

@@ -1,7 +1,8 @@
 {
     "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20220324.json",
-    "version": "0.4.8",
+    "version": "0.4.9",
     "changelog": {
+        "0.4.9": "fix orientation issue on clicks",
         "0.4.8": "Add infer transforms to manage clicks from viewer",
         "0.4.7": "fix the wrong GPU index issue of multi-node",
         "0.4.6": "update to use rc7 which solves dynunet issue",
@@ -31,7 +32,8 @@
     "optional_packages_version": {
         "itk": "5.3.0",
         "pytorch-ignite": "0.4.9",
-        "scikit-image": "0.19.3"
+        "scikit-image": "0.19.3",
+        "einops": "0.6.1"
     },
     "name": "Spleen DeepEdit annotation",
     "task": "Decathlon spleen segmentation",

models/spleen_deepedit_annotation/scripts/transforms.py

@@ -0,0 +1,38 @@
+from typing import Dict
+
+import numpy as np
+from einops import rearrange
+from monai.transforms.transform import Transform
+
+
+class OrientationGuidanceMultipleLabelDeepEditd(Transform):
+    def __init__(self, ref_image="image", label_names=None):
+        """
+        Convert the guidance to the RAS orientation
+        """
+        self.ref_image = ref_image
+        self.label_names = label_names
+
+    def transform_points(self, point, affine):
+        """transform point to the coordinates of the transformed image
+        point: numpy array [bs, N, 3]
+        """
+        bs, N = point.shape[:2]
+        point = np.concatenate((point, np.ones((bs, N, 1))), axis=-1)
+        point = rearrange(point, "b n d -> d (b n)")
+        point = affine @ point
+        point = rearrange(point, "d (b n)-> b n d", b=bs)[:, :, :3]
+        return point
+
+    def __call__(self, data):
+        d: Dict = dict(data)
+        for key_label in self.label_names.keys():
+            points = d.get(key_label, [])
+            if len(points) < 1:
+                continue
+            reoriented_points = self.transform_points(
+                np.array(points)[None],
+                np.linalg.inv(d[self.ref_image].meta["affine"].numpy()) @ d[self.ref_image].meta["original_affine"],
+            )
+            d[key_label] = reoriented_points[0]
+        return d