update wholebody seg

Signed-off-by: Yiheng Wang <[email protected]>
Project-MONAI · Feb 22, 2024 · 4f41bba · 4f41bba
1 parent 7059a9e
commit 4f41bba
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 20 deletions.
diff --git a/models/wholeBody_ct_segmentation/configs/inference.json b/models/wholeBody_ct_segmentation/configs/inference.json
@@ -6,10 +6,16 @@
     },
     "imports": [
         "$import glob",
+        "$import numpy",
         "$import os"
     ],
     "bundle_root": ".",
+    "image_key": "image",
     "output_dir": "$@bundle_root + '/eval'",
+    "output_ext": ".nii.gz",
+    "output_dtype": "$numpy.float32",
+    "output_postfix": "trans",
+    "separate_folder": true,
     "dataset_dir": "sampledata",
     "datalist": "$list(sorted(glob.glob(@dataset_dir + '/imagesTs/*.nii.gz')))",
     "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')",
@@ -40,35 +46,35 @@
         "transforms": [
             {
                 "_target_": "LoadImaged",
-                "keys": "image"
+                "keys": "@image_key"
             },
             {
                 "_target_": "EnsureTyped",
-                "keys": "image"
+                "keys": "@image_key"
             },
             {
                 "_target_": "EnsureChannelFirstd",
-                "keys": "image"
+                "keys": "@image_key"
             },
             {
                 "_target_": "Orientationd",
-                "keys": "image",
+                "keys": "@image_key",
                 "axcodes": "RAS"
             },
             {
                 "_target_": "Spacingd",
-                "keys": "image",
+                "keys": "@image_key",
                 "pixdim": "@pixdim",
                 "mode": "bilinear"
             },
             {
                 "_target_": "NormalizeIntensityd",
-                "keys": "image",
+                "keys": "@image_key",
                 "nonzero": true
             },
             {
                 "_target_": "ScaleIntensityd",
-                "keys": "image",
+                "keys": "@image_key",
                 "minv": -1.0,
                 "maxv": 1.0
             }
@@ -116,25 +122,22 @@
                 "_target_": "Invertd",
                 "keys": "pred",
                 "transform": "@preprocessing",
-                "orig_keys": "image",
+                "orig_keys": "@image_key",
                 "nearest_interp": true,
                 "to_tensor": true
             },
             {
                 "_target_": "SaveImaged",
                 "keys": "pred",
-                "output_dir": "@output_dir"
+                "output_dir": "@output_dir",
+                "output_ext": "@output_ext",
+                "output_dtype": "@output_dtype",
+                "output_postfix": "@output_postfix",
+                "separate_folder": "@separate_folder"
             }
         ]
     },
     "handlers": [
-        {
-            "_target_": "CheckpointLoader",
-            "load_path": "$@bundle_root + '/models/' + @modelname",
-            "load_dict": {
-                "model": "@network"
-            }
-        },
         {
             "_target_": "StatsHandler",
             "iteration_log": false
@@ -150,8 +153,16 @@
         "val_handlers": "@handlers",
         "amp": true
     },
+    "checkpointloader": {
+        "_target_": "CheckpointLoader",
+        "load_path": "$@bundle_root + '/models/' + @modelname",
+        "load_dict": {
+            "model": "@network"
+        }
+    },
     "initialize": [
-        "$setattr(torch.backends.cudnn, 'benchmark', True)"
+        "$setattr(torch.backends.cudnn, 'benchmark', True)",
+        "$@checkpointloader(@evaluator)"
     ],
     "run": [
         "[email protected]()"

diff --git a/models/wholeBody_ct_segmentation/configs/inference_trt.json b/models/wholeBody_ct_segmentation/configs/inference_trt.json
@@ -4,7 +4,9 @@
         "$import os",
         "$import torch_tensorrt"
     ],
-    "handlers#0#_disabled_": true,
     "network_def": "$torch.jit.load(@bundle_root + '/models/model_trt.ts')",
-    "evaluator#amp": false
+    "evaluator#amp": false,
+    "initialize": [
+        "$setattr(torch.backends.cudnn, 'benchmark', True)"
+    ]
 }
diff --git a/models/wholeBody_ct_segmentation/configs/metadata.json b/models/wholeBody_ct_segmentation/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.2.1",
+    "version": "0.2.2",
     "changelog": {
+        "0.2.2": "add checkpoint loader for infer",
         "0.2.1": "remove meta_dict usage",
         "0.2.0": "add support for TensorRT conversion and inference",
         "0.1.9": "fix the wrong GPU index issue of multi-node",