Load Config Options from Bundle (#1302)

* add bundle configs option support

Signed-off-by: tangy5 <[email protected]>

* add bundle configs option support

Signed-off-by: tangy5 <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update with merged bundle

Signed-off-by: tangy5 <[email protected]>

* small fix

Signed-off-by: tangy5 <[email protected]>

* small fix

Signed-off-by: tangy5 <[email protected]>

* small fix

Signed-off-by: tangy5 <[email protected]>

* sanity checks

Signed-off-by: tangy5 <[email protected]>

* update

Signed-off-by: tangy5 <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: tangy5 <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: tangy5

monailabel/tasks/infer/basic_infer.py

@@ -267,6 +267,11 @@ def __call__(
         req = copy.deepcopy(self._config)
         req.update(request)
 
+        # model options
+        self.path.append(
+            os.path.join(os.path.dirname(self.path[0]), req.get("model_filename", "model.pt"))
+        ) if self.path and isinstance(self.path, list) else self.path
+
         # device
         device = req.get("device", "cuda")
         device = device if isinstance(device, str) else device[0]

monailabel/tasks/infer/bundle.py

@@ -9,6 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import glob
 import json
 import logging
 import os
@@ -66,6 +67,9 @@ def key_detector(self) -> Sequence[str]:
     def key_detector_ops(self) -> Sequence[str]:
         return ["detector_ops"]
 
+    def key_displayable_configs(self) -> Sequence[str]:
+        return ["displayable_configs"]
+
 
 class BundleInferTask(BasicInferTask):
     """
@@ -147,6 +151,17 @@ def __init__(
             preload=strtobool(conf.get("preload", "false")),
             **kwargs,
         )
+
+        # Add models options if more than one model is provided by bundle.
+        pytorch_models = [os.path.basename(p) for p in glob.glob(os.path.join(path, "models", "*.pt"))]
+        pytorch_models.sort(key=len)
+        self._config.update({"model_filename": pytorch_models})
+        # Add bundle's loadable params to MONAI Label config, load exposed keys and params to options panel
+        for k in self.const.key_displayable_configs():
+            if self.bundle_config.get(k):
+                self.displayable_configs = self.bundle_config.get_parsed_content(k, instantiate=True)  # type: ignore
+                self._config.update(self.displayable_configs)
+
         self.valid = True
         self.version = metadata.get("version")
         sys.path.remove(self.bundle_path)
@@ -160,6 +175,12 @@ def info(self) -> Dict[str, Any]:
         return i
 
     def pre_transforms(self, data=None) -> Sequence[Callable]:
+        # Update bundle parameters based on user's option
+        for k in self.const.key_displayable_configs():
+            if self.bundle_config.get(k):
+                self.bundle_config[k].update({c: data[c] for c in self.displayable_configs.keys()})
+                self.bundle_config.parse()
+
         sys.path.insert(0, self.bundle_path)
         unload_module("scripts")
         self._update_device(data)

monailabel/tasks/train/bundle.py

@@ -9,6 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import glob
 import json
 import logging
 import os
@@ -79,6 +80,9 @@ def key_experiment_name(self) -> str:
     def key_run_name(self) -> str:
         return "run_name"
 
+    def key_displayable_configs(self) -> Sequence[str]:
+        return ["displayable_configs"]
+
 
 class BundleTrainTask(TrainTask):
     def __init__(
@@ -127,7 +131,11 @@ def info(self):
         return i
 
     def config(self):
-        return {
+        # Add models and param optiom to train option panel
+        pytorch_models = [os.path.basename(p) for p in glob.glob(os.path.join(self.bundle_path, "models", "*.pt"))]
+        pytorch_models.sort(key=len)
+
+        config_options = {
             "device": "cuda",  # DEVICE
             "pretrained": True,  # USE EXISTING CHECKPOINT/PRETRAINED MODEL
             "max_epochs": 50,  # TOTAL EPOCHS TO RUN
@@ -139,8 +147,15 @@ def config(self):
             else ["None", "mlflow"],
             "tracking_uri": settings.MONAI_LABEL_TRACKING_URI,
             "tracking_experiment_name": "",
+            "model_filename": pytorch_models,
         }
 
+        for k in self.const.key_displayable_configs():
+            if self.bundle_config.get(k):
+                config_options.update(self.bundle_config.get_parsed_content(k, instantiate=True))  # type: ignore
+
+        return config_options
+
     def _fetch_datalist(self, request, datastore: Datastore):
         datalist = datastore.datalist()
 
@@ -177,8 +192,8 @@ def _partition_datalist(self, datalist, request, shuffle=False):
         logger.info(f"Total Records for Validation: {len(val_datalist) if val_datalist else ''}")
         return train_datalist, val_datalist
 
-    def _load_checkpoint(self, output_dir, pretrained, train_handlers):
-        load_path = os.path.join(output_dir, self.const.model_pytorch()) if pretrained else None
+    def _load_checkpoint(self, model_pytorch, pretrained, train_handlers):
+        load_path = model_pytorch if pretrained else None
         if os.path.exists(load_path):
             logger.info(f"Add Checkpoint Loader for Path: {load_path}")
 
@@ -226,7 +241,9 @@ def __call__(self, request, datastore: Datastore):
         logger.info(f"(Experiment Management) Run Name: {tracking_run_name}")
 
         train_handlers = self.bundle_config.get(self.const.key_train_handlers(), [])
-        self._load_checkpoint(os.path.join(self.bundle_path, "models"), pretrained, train_handlers)
+
+        model_pytorch = os.path.join(self.bundle_path, "models", request.get("model_filename", "model.pt"))
+        self._load_checkpoint(model_pytorch, pretrained, train_handlers)
 
         overrides = {
             self.const.key_bundle_root(): self.bundle_path,
@@ -236,6 +253,12 @@ def __call__(self, request, datastore: Datastore):
             self.const.key_train_handlers(): train_handlers,
         }
 
+        # update config options from user
+        for k in self.const.key_displayable_configs():
+            if self.bundle_config.get(k):
+                displayable_configs = self.bundle_config.get_parsed_content(k, instantiate=True)
+                overrides[k] = {c: request[c] for c in displayable_configs.keys()}
+
         if tracking and tracking.lower() != "none":
             overrides[self.const.key_tracking()] = tracking
             if tracking_uri:

monailabel/utils/others/modelzoo_list.py

@@ -31,4 +31,5 @@
     "swin_unetr_btcv_segmentation",  # 3D transformer model for multi-organ segmentation. Added Oct 2022
     "wholeBrainSeg_Large_UNEST_segmentation",  # whole brain segmentation for T1 MRI brain images. Added Oct 2022
     "lung_nodule_ct_detection",  # The first lung nodule detection task can be used for MONAI Label. Added Dec 2022
+    "wholeBody_ct_segmentation",  # The SegResNet trained TotalSegmentator dataset with 104 tissues. Added Feb 2023
 ]

sample-apps/monaibundle/README.md

@@ -46,6 +46,7 @@ The Bundle App supports most labeling models in the Model Zoo, please see the ta
 | [wholeBrainSeg_UNEST_segmentation](https://github.com/Project-MONAI/model-zoo/tree/dev/models/wholeBrainSeg_Large_UNEST_segmentation) | UNesT | Whole Brain | MRI T1 |  A pre-trained for inference (3D) 133 whole brain structures segmentation |
 | [spleen_deepedit_annotation](https://github.com/Project-MONAI/model-zoo/tree/dev/models/spleen_deepedit_annotation) | DeepEdit | Spleen| CT | An interactive method for 3D spleen Segmentation |
 | [lung_nodule_ct_detection](https://github.com/Project-MONAI/model-zoo/tree/dev/models/lung_nodule_ct_detection) | RetinaNet | Lung Nodule| CT | The detection model for 3D CT images |
+| [wholeBody_ct_segmentation](https://github.com/Project-MONAI/model-zoo/tree/dev/models/wholeBody_ct_segmentation) | SegResNet | 104 body structures| CT | The segmentation model for 104 tissue from 3D CT images (TotalSegmentator Dataset) |
 
 Supported tasks update based on [Model-Zoo](https://github.com/Project-MONAI/model-zoo/tree/dev/models) release.