Update DeepEdit infer to accept clicks from viewer (Project-MONAI#491)

Update DeepEdit Infer config to accept clicks from the user --------- Signed-off-by: Andres <diazandr3s@gmail.com> Signed-off-by: Yiheng Wang <vennw@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Yiheng Wang <vennw@nvidia.com>
yiheng-wang-nv · Aug 28, 2023 · 85b8cf2 · 85b8cf2
1 parent 7117a85
commit 85b8cf2
Show file tree

Hide file tree

Showing 6 changed files with 114 additions and 44 deletions.
diff --git a/.github/workflows/premerge-cpu.yml b/.github/workflows/premerge-cpu.yml
@@ -40,5 +40,6 @@ jobs:
       run: |
         # clean up temporary files
         $(pwd)/runtests.sh --clean
+        df -h
         bash $(pwd)/ci/run_premerge_cpu.sh changed
       shell: bash
diff --git a/ci/run_premerge_cpu.sh b/ci/run_premerge_cpu.sh
@@ -43,7 +43,7 @@ remove_pipenv() {
 
 verify_bundle() {
     echo 'Run verify bundle...'
-    init_pipenv requirements-dev.txt
+    init_pipenv requirements.txt
     head_ref=$(git rev-parse HEAD)
     git fetch origin dev $head_ref
     # achieve all changed files in 'models'

diff --git a/ci/unit_tests/test_spleen_deepedit_annotation.py b/ci/unit_tests/test_spleen_deepedit_annotation.py
@@ -120,6 +120,25 @@ def test_infer_config(self, override):
         )
         check_workflow(inferrer, check_properties=True)
 
+    @parameterized.expand([TEST_CASE_2])
+    def test_infer_click_config(self, override):
+        override["dataset_dir"] = self.dataset_dir
+        override["use_click"] = True
+        override[
+            "dataset#data"
+        ] = "$[{'image': i, 'background': [], 'spleen': [[6, 6, 6], [8, 8, 8]]} for i in @datalist]"
+        bundle_root = override["bundle_root"]
+        print(override)
+
+        inferrer = ConfigWorkflow(
+            workflow="infer",
+            config_file=os.path.join(bundle_root, "configs/inference.json"),
+            logging_file=os.path.join(bundle_root, "configs/logging.conf"),
+            meta_file=os.path.join(bundle_root, "configs/metadata.json"),
+            **override,
+        )
+        check_workflow(inferrer, check_properties=True)
+
 
 if __name__ == "__main__":
     loader = unittest.TestLoader()

diff --git a/models/spleen_deepedit_annotation/configs/inference.json b/models/spleen_deepedit_annotation/configs/inference.json
@@ -19,6 +19,7 @@
     ],
     "number_intensity_ch": 1,
     "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')",
+    "use_click": false,
     "network_def": {
         "_target_": "DynUNet",
         "spatial_dims": 3,
@@ -60,49 +61,79 @@
         "res_block": true
     },
     "network": "$@network_def.to(@device)",
+    "preprocessing_transforms": [
+        {
+            "_target_": "LoadImaged",
+            "keys": "image",
+            "reader": "ITKReader"
+        },
+        {
+            "_target_": "EnsureChannelFirstd",
+            "keys": "image"
+        },
+        {
+            "_target_": "Orientationd",
+            "keys": "image",
+            "axcodes": "RAS"
+        },
+        {
+            "_target_": "ScaleIntensityRanged",
+            "keys": "image",
+            "a_min": -175,
+            "a_max": 250,
+            "b_min": 0.0,
+            "b_max": 1.0,
+            "clip": true
+        }
+    ],
+    "auto_seg_transforms": [
+        {
+            "_target_": "Resized",
+            "keys": "image",
+            "spatial_size": "@spatial_size",
+            "mode": "area"
+        },
+        {
+            "_target_": "DiscardAddGuidanced",
+            "keys": "image",
+            "label_names": "@label_names",
+            "number_intensity_ch": "@number_intensity_ch"
+        }
+    ],
+    "deepedit_transforms": [
+        {
+            "_target_": "AddGuidanceFromPointsDeepEditd",
+            "ref_image": "image",
+            "guidance": "guidance",
+            "label_names": "@label_names"
+        },
+        {
+            "_target_": "Resized",
+            "keys": "image",
+            "spatial_size": "@spatial_size",
+            "mode": "area"
+        },
+        {
+            "_target_": "ResizeGuidanceMultipleLabelDeepEditd",
+            "guidance": "guidance",
+            "ref_image": "image"
+        },
+        {
+            "_target_": "AddGuidanceSignalDeepEditd",
+            "keys": "image",
+            "guidance": "guidance",
+            "number_intensity_ch": "@number_intensity_ch"
+        }
+    ],
+    "extra_transforms": [
+        {
+            "_target_": "EnsureTyped",
+            "keys": "image"
+        }
+    ],
     "preprocessing": {
         "_target_": "Compose",
-        "transforms": [
-            {
-                "_target_": "LoadImaged",
-                "keys": "image",
-                "reader": "ITKReader"
-            },
-            {
-                "_target_": "EnsureChannelFirstd",
-                "keys": "image"
-            },
-            {
-                "_target_": "Orientationd",
-                "keys": "image",
-                "axcodes": "RAS"
-            },
-            {
-                "_target_": "ScaleIntensityRanged",
-                "keys": "image",
-                "a_min": -175,
-                "a_max": 250,
-                "b_min": 0.0,
-                "b_max": 1.0,
-                "clip": true
-            },
-            {
-                "_target_": "Resized",
-                "keys": "image",
-                "spatial_size": "@spatial_size",
-                "mode": "area"
-            },
-            {
-                "_target_": "DiscardAddGuidanced",
-                "keys": "image",
-                "label_names": "@label_names",
-                "number_intensity_ch": "@number_intensity_ch"
-            },
-            {
-                "_target_": "EnsureTyped",
-                "keys": "image"
-            }
-        ]
+        "transforms": "$@preprocessing_transforms + (@deepedit_transforms if @use_click else @auto_seg_transforms) + @extra_transforms"
     },
     "dataset": {
         "_target_": "Dataset",

diff --git a/models/spleen_deepedit_annotation/configs/metadata.json b/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.7",
+    "version": "0.4.8",
     "changelog": {
+        "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",
         "0.4.5": "remove error dollar symbol in readme",
@@ -28,7 +29,7 @@
     "pytorch_version": "1.13.1",
     "numpy_version": "1.22.2",
     "optional_packages_version": {
-        "nibabel": "4.0.1",
+        "itk": "5.3.0",
         "pytorch-ignite": "0.4.9",
         "scikit-image": "0.19.3"
     },

diff --git a/models/spleen_deepedit_annotation/docs/README.md b/models/spleen_deepedit_annotation/docs/README.md
@@ -118,10 +118,28 @@ python -m monai.bundle run --config_file "['configs/train.json','configs/evaluat
 
 ####  Execute inference:
 
+
+For automatic inference mode:
+
+
 ```
 python -m monai.bundle run --config_file configs/inference.json
 ```
 
+For interactive segmentation mode, in which the user provides clicks, set the **use_click** flag to true:
+
+
+```
+python -m monai.bundle run --config_file configs/inference.json --use_click true
+```
+
+Clicks should be added to the data dictionary that is passed to the preprocessing transforms. The add keys are defined in `label_names` in `configs/inference.json`, and the corresponding values are the point coordinates. The following is an example of a data dictionary:
+
+```
+{"image": "example.nii.gz", "background": [], "spleen": [[I1, J1, K1], [I2, J2, K2]]}
+```
+where **[I1,J1,K1]** and **[I2,J2,K2]** are the point coordinates.
+
 #### Export checkpoint to TensorRT based models with fp32 or fp16 precision:
 
 ```bash