558 Update checkpoint loader for more bundles (Project-MONAI#561)

Part of Project-MONAI#558  .

### Description
A few sentences describing the changes proposed in this pull request.

### Status
**Ready/Work in progress/Hold**

### Please ensure all the checkboxes:
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Codeformat tests passed locally by running `./runtests.sh
--codeformat`.
- [ ] In-line docstrings updated.
- [ ] Update `version` and `changelog` in `metadata.json` if changing an
existing bundle.
- [ ] Please ensure the naming rules in config files meet our
requirements (please refer to: `CONTRIBUTING.md`).
- [ ] Ensure versions of packages such as `monai`, `pytorch` and `numpy`
are correct in `metadata.json`.
- [ ] Descriptions should be consistent with the content, such as
`eval_metrics` of the provided weights and TorchScript modules.
- [ ] Files larger than 25MB are excluded and replaced by providing
download links in `large_file.yml`.
- [ ] Avoid using path that contains personal information within config
files (such as use `/home/your_name/` for `"bundle_root"`).

---------

Signed-off-by: Yiheng Wang <[email protected]>

ci/unit_tests/test_endoscopic_tool_segmentation.py

@@ -31,14 +31,23 @@
 ]
 
 TEST_CASE_2 = [  # inference
-    {
-        "bundle_root": "models/endoscopic_tool_segmentation",
-        "handlers#0#_disabled_": True,
-        "preprocessing#transforms#2#spatial_size": [32, 32],
-    }
+    {"bundle_root": "models/endoscopic_tool_segmentation", "preprocessing#transforms#2#spatial_size": [32, 32]}
 ]
 
 
+def test_order(test_name1, test_name2):
+    def get_order(name):
+        if "train" in name:
+            return 1
+        if "eval" in name:
+            return 2
+        if "infer" in name:
+            return 3
+        return 4
+
+    return get_order(test_name1) - get_order(test_name2)
+
+
 class TestEndoscopicSeg(unittest.TestCase):
     def setUp(self):
         self.dataset_dir = tempfile.mkdtemp()
@@ -105,4 +114,6 @@ def test_infer_config(self, override):
 
 
 if __name__ == "__main__":
-    unittest.main()
+    loader = unittest.TestLoader()
+    loader.sortTestMethodsUsing = test_order
+    unittest.main(testLoader=loader)

models/endoscopic_tool_segmentation/configs/inference.json

@@ -1,10 +1,16 @@
 {
     "imports": [
         "$import glob",
+        "$import numpy",
         "$import os"
     ],
     "bundle_root": ".",
+    "image_key": "image",
     "output_dir": "$@bundle_root + '/eval'",
+    "output_ext": ".png",
+    "output_dtype": "$numpy.float32",
+    "output_postfix": "trans",
+    "separate_folder": true,
     "dataset_dir": "/workspace/data/endoscopic_tool_dataset",
     "datalist": "$list(sorted(glob.glob(os.path.join(@dataset_dir,'test', '*', '*[!seg].jpg'))))",
     "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')",
@@ -24,18 +30,18 @@
         "transforms": [
             {
                 "_target_": "LoadImaged",
-                "keys": "image",
+                "keys": "@image_key",
                 "image_only": false
             },
             {
                 "_target_": "EnsureChannelFirstd",
-                "keys": "image",
+                "keys": "@image_key",
                 "channel_dim": -1
             },
             {
                 "_target_": "Resized",
                 "keys": [
-                    "image"
+                    "@image_key"
                 ],
                 "spatial_size": [
                     736,
@@ -48,7 +54,7 @@
             {
                 "_target_": "ScaleIntensityd",
                 "keys": [
-                    "image"
+                    "@image_key"
                 ]
             }
         ]
@@ -75,7 +81,7 @@
                 "_target_": "Invertd",
                 "keys": "pred",
                 "transform": "@preprocessing",
-                "orig_keys": "image",
+                "orig_keys": "@image_key",
                 "meta_key_postfix": "meta_dict",
                 "nearest_interp": false,
                 "to_tensor": true
@@ -100,21 +106,16 @@
                 "keys": "pred",
                 "meta_keys": "pred_meta_dict",
                 "output_dir": "@output_dir",
-                "output_ext": ".png",
+                "output_ext": "@output_ext",
+                "output_dtype": "@output_dtype",
+                "output_postfix": "@output_postfix",
+                "separate_folder": "@separate_folder",
                 "scale": 255,
                 "squeeze_end_dims": true
             }
         ]
     },
     "handlers": [
-        {
-            "_target_": "CheckpointLoader",
-            "load_path": "$@bundle_root + '/models/model.pt'",
-            "load_dict": {
-                "model": "@network"
-            },
-            "map_location": "@device"
-        },
         {
             "_target_": "StatsHandler",
             "iteration_log": false
@@ -129,8 +130,17 @@
         "postprocessing": "@postprocessing",
         "val_handlers": "@handlers"
     },
+    "checkpointloader": {
+        "_target_": "CheckpointLoader",
+        "load_path": "$@bundle_root + '/models/model.pt'",
+        "load_dict": {
+            "model": "@network"
+        },
+        "map_location": "@device"
+    },
     "initialize": [
-        "$monai.utils.set_determinism(seed=123)"
+        "$monai.utils.set_determinism(seed=123)",
+        "$@checkpointloader(@evaluator)"
     ],
     "run": [
         "[email protected]()"

models/endoscopic_tool_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": [
+        "$monai.utils.set_determinism(seed=123)"
+    ]
 }

models/endoscopic_tool_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.5.6",
+    "version": "0.5.7",
     "changelog": {
+        "0.5.7": "add checkpoint loader for infer",
         "0.5.6": "update to use monai 1.3.0",
         "0.5.5": "update AddChanneld with EnsureChannelFirstd and set image_only to False",
         "0.5.4": "fix the wrong GPU index issue of multi-node",

models/pancreas_ct_dints_segmentation/configs/inference.yaml

@@ -1,13 +1,19 @@
 ---
 imports:
 - "$import glob"
+- "$import numpy"
 - "$import os"
 input_channels: 1
 output_classes: 3
 arch_ckpt_path: "$@bundle_root + '/models/search_code_18590.pt'"
 arch_ckpt: "$torch.load(@arch_ckpt_path, map_location=torch.device('cuda'))"
 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: "/workspace/data/msd/Task07_Pancreas"
 data_list_file_path: "$@bundle_root + '/configs/dataset_0.json'"
 datalist: "$monai.data.load_decathlon_datalist(@data_list_file_path, data_list_key='testing',
@@ -35,28 +41,28 @@ preprocessing:
   _target_: Compose
   transforms:
   - _target_: LoadImaged
-    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:
     - 1
     - 1
     - 1
     mode: bilinear
   - _target_: ScaleIntensityRanged
-    keys: image
+    keys: "@image_key"
     a_min: -87
     a_max: 199
     b_min: 0
     b_max: 1
     clip: true
   - _target_: EnsureTyped
-    keys: image
+    keys: "@image_key"
 dataset:
   _target_: Dataset
   data: "@datalist"
@@ -84,7 +90,7 @@ postprocessing:
   - _target_: Invertd
     keys: pred
     transform: "@preprocessing"
-    orig_keys: image
+    orig_keys: "@image_key"
     meta_key_postfix: meta_dict
     nearest_interp: false
     to_tensor: true
@@ -95,11 +101,11 @@ postprocessing:
     keys: pred
     meta_keys: pred_meta_dict
     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/model.pt'"
-  load_dict:
-    model: "@network"
 - _target_: StatsHandler
   iteration_log: false
 evaluator:
@@ -111,7 +117,13 @@ evaluator:
   postprocessing: "@postprocessing"
   val_handlers: "@handlers"
   amp: true
+checkpointloader:
+  _target_: CheckpointLoader
+  load_path: "$@bundle_root + '/models/model.pt'"
+  load_dict:
+    model: "@network"
 initialize:
 - "$setattr(torch.backends.cudnn, 'benchmark', True)"
+- "$@checkpointloader(@evaluator)"
 run:
 - "[email protected]()"

models/pancreas_ct_dints_segmentation/configs/inference_trt.yaml

@@ -3,6 +3,7 @@ imports:
 - "$import glob"
 - "$import os"
 - "$import torch_tensorrt"
-handlers#0#_disabled_: true
 network_def: "$torch.jit.load(@bundle_root + '/models/model_trt.ts')"
 evaluator#amp: false
+initialize:
+- "$setattr(torch.backends.cudnn, 'benchmark', True)"

models/pancreas_ct_dints_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.4.5",
+    "version": "0.4.6",
     "changelog": {
+        "0.4.6": "add checkpoint loader for infer",
         "0.4.5": "set image_only to False",
         "0.4.4": "update the benchmark results of TensorRT",
         "0.4.3": "add support for TensorRT conversion and inference",

models/pathology_nuclick_annotation/configs/inference.json

@@ -1,12 +1,18 @@
 {
     "imports": [
         "$import glob",
+        "$import numpy",
         "$import json",
         "$import pathlib",
         "$import os"
     ],
     "bundle_root": ".",
+    "image_key": "image",
     "output_dir": "$@bundle_root + '/eval'",
+    "output_ext": ".png",
+    "output_dtype": "$numpy.float32",
+    "output_postfix": "trans",
+    "separate_folder": true,
     "dataset_dir": "/workspace/data/CoNSePNuclei",
     "images": "$list(sorted(glob.glob(@dataset_dir + '/Test/Images/*.png')))[:1]",
     "centroids": "$list(sorted(glob.glob(@dataset_dir + '/Test/Centroids/*.txt')))[:1]",
@@ -32,31 +38,31 @@
         "transforms": [
             {
                 "_target_": "LoadImaged",
-                "keys": "image",
+                "keys": "@image_key",
                 "dtype": "uint8",
                 "image_only": false
             },
             {
                 "_target_": "EnsureChannelFirstd",
-                "keys": "image"
+                "keys": "@image_key"
             },
             {
                 "_target_": "ScaleIntensityRanged",
-                "keys": "image",
+                "keys": "@image_key",
                 "a_min": 0.0,
                 "a_max": 255.0,
                 "b_min": -1.0,
                 "b_max": 1.0
             },
             {
                 "_target_": "AddClickSignalsd",
-                "image": "image",
+                "image": "@image_key",
                 "foreground": "foreground",
                 "gaussian": false
             },
             {
                 "_target_": "SqueezeDimd",
-                "keys": "image"
+                "keys": "@image_key"
             }
         ]
     },
@@ -97,18 +103,14 @@
                 "keys": "pred",
                 "meta_keys": "pred_meta_dict",
                 "output_dir": "@output_dir",
-                "output_ext": ".png"
+                "output_ext": "@output_ext",
+                "output_dtype": "@output_dtype",
+                "output_postfix": "@output_postfix",
+                "separate_folder": "@separate_folder"
             }
         ]
     },
     "handlers": [
-        {
-            "_target_": "CheckpointLoader",
-            "load_path": "$@bundle_root + '/models/model.pt'",
-            "load_dict": {
-                "model": "@network"
-            }
-        },
         {
             "_target_": "StatsHandler",
             "iteration_log": false
@@ -124,10 +126,18 @@
         "val_handlers": "@handlers",
         "amp": true
     },
+    "checkpointloader": {
+        "_target_": "CheckpointLoader",
+        "load_path": "$@bundle_root + '/models/model.pt'",
+        "load_dict": {
+            "model": "@network"
+        }
+    },
     "initialize": [
         "$import sys",
         "$sys.path.append(@bundle_root)",
-        "$monai.utils.set_determinism(seed=123)"
+        "$monai.utils.set_determinism(seed=123)",
+        "$@checkpointloader(@evaluator)"
     ],
     "run": [
         "[email protected]()"

models/pathology_nuclick_annotation/configs/inference_trt.json

@@ -6,7 +6,11 @@
         "$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": [
+        "$import sys",
+        "$sys.path.append(@bundle_root)",
+        "$monai.utils.set_determinism(seed=123)"
+    ]
 }

models/pathology_nuclick_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.1.6",
+    "version": "0.1.7",
     "changelog": {
+        "0.1.7": "add checkpoint loader for infer",
         "0.1.6": "set image_only to False",
         "0.1.5": "add support for TensorRT conversion and inference",
         "0.1.4": "fix the wrong GPU index issue of multi-node",