Epistemic Scoring for monaibundle app (#1109)

Author: tangy5

monailabel/utils/others/generic.py

@@ -247,16 +247,17 @@ def get_bundle_models(app_dir, conf, conf_key="models"):
 
     model_dir = os.path.join(app_dir, "model")
 
-    zoo_info = get_all_bundles_list()
     zoo_source = conf.get("zoo_source", MONAI_ZOO_SOURCE)
     zoo_repo = conf.get("zoo_repo", MONAI_ZOO_REPO)
+    auth_token = conf.get("auth_token", None)
+    zoo_info = get_all_bundles_list(auth_token=auth_token)
 
     # filter model zoo bundle with MONAI Label supported bundles according to the maintaining list, return all version bundles list
     available = [
         i[0] + "_v" + v
         for i in zoo_info
         if i[0] in MAINTAINED_BUNDLES
-        for v in get_bundle_versions(i[0])["all_versions"]
+        for v in get_bundle_versions(i[0], auth_token=auth_token)["all_versions"]
     ]
 
     models = conf.get(conf_key)

sample-apps/monaibundle/README.md

@@ -50,6 +50,12 @@ Supported tasks update based on [Model-Zoo](https://github.com/Project-MONAI/mod
 
 Note: MONAI Label support labeling bundle models in the MODEL ZOO, non-labeling tasks are not available for MONAI Label.
 
+Note: the monaibundle app uses monai bundle api to get information of latest Model-Zoo. In order to increase the rate limits of calling GIthub APIs, you can input your personal access token.
+    Please check the following link for more details about rate limiting:
+    https://docs.github.com/en/rest/overview/resources-in-the-rest-api#rate-limiting
+
+Set **--conf auth_token <github personal access token>** to increate rate limits.
+
 ### How To Use?
 
 ```bash
@@ -72,14 +78,38 @@ monailabel start_server --app workspace/monaibundle --studies workspace/images -
 monailabel start_server --app workspace/monaibundle --studies workspace/images --conf models spleen_ct_segmentation_v0.1.0 --conf skip_trainers true
 ```
 
+### Epistemic Scoring for monaibundle app
+
+The monaibundle app supports epistemic scoring using bundles models. The monaibundle consumes **epistemic_model** as config parameter.
+If a valid scoring bundle model is provided with **--conf epistemic_model <bundlename>**, scoring inference will be triggered.
+The loaded scoring bundle model can be either model-zoo models or local bundles, the network must support **dropout** argument.
+
+```bash
+# Use the UNet in spleen_ct_segmentation_v0.2.0 bundle as epistemic scoring model.
+# Manual define epistemic scoring parameters
+monailabel start_server \
+  --app workspace/monaibundle \
+  --studies workspace/images \
+  --conf models spleen_ct_segmentation_v0.2.0,swin_unetr_btcv_segmentation_v0.2.0 \
+  --conf epistemic_model spleen_ct_segmentation_v0.2.0
+  --conf epistemic_max_samples 0 \
+  --conf epistemic_simulation_size 5
+  --conf epistemic_dropout 0.2
+
+```
+
+Users can then select active learning strategies for selecting data to label.
+
 #### Additional Configs
 
 Pass them as **--conf _name_ _value_** while starting MONAILabelServer
 
-| Name          | Values          | Description                                                                                 |
-|---------------|-----------------|---------------------------------------------------------------------------------------------|
-| zoo_info      | string          | _Default value:_ https://github.com/Project-MONAI/model-zoo/blob/dev/models/model_info.json |
-| zoo_source    | string          | _Default value:_ github                                                                     |
-| zoo_repo      | string          | _Default value:_ Project-MONAI/model-zoo/hosting_storage_v1                                 |
-| preload       | true, **false** | Preload model into GPU                                                                      |
-| skip_trainers | true, **false** | Skip adding training tasks (Run in Infer mode only)                                         |
+| Name                      | Values          | Description                                                                                 |
+|---------------------------|-----------------|---------------------------------------------------------------------------------------------|
+| zoo_source                | string          | _Default value:_ github                                                                     |
+| zoo_repo                  | string          | _Default value:_ Project-MONAI/model-zoo/hosting_storage_v1                                 |
+| preload                   | true, **false** | Preload model into GPU                                                                      |
+| skip_trainers             | true, **false** | Skip adding training tasks (Run in Infer mode only)                                         |
+| epistemic_max_samples     | int             | _Default value:_ 0    ;  Epistemic scoring parameters                                       |
+| epistemic_simulation_size | int             | _Default value:_ 5    ;  Epistemic simulation size parameters                               |
+| epistemic_dropout         | float           | _Default value:_ 0.2  ;  Epistemic scoring parameters: Dropout rate for scoring models      |                               |

sample-apps/monaibundle/main.py

@@ -13,14 +13,18 @@
 import os
 from typing import Dict
 
+from monai.transforms import Invertd, SaveImaged
+
 import monailabel
 from monailabel.interfaces.app import MONAILabelApp
 from monailabel.interfaces.tasks.infer_v2 import InferTask
+from monailabel.interfaces.tasks.scoring import ScoringMethod
 from monailabel.interfaces.tasks.strategy import Strategy
 from monailabel.interfaces.tasks.train import TrainTask
 from monailabel.tasks.activelearning.first import First
 from monailabel.tasks.activelearning.random import Random
 from monailabel.tasks.infer.bundle import BundleInferTask
+from monailabel.tasks.scoring.epistemic_v2 import EpistemicScoring
 from monailabel.tasks.train.bundle import BundleTrainTask
 from monailabel.utils.others.generic import get_bundle_models, strtobool
 
@@ -30,6 +34,15 @@
 class MyApp(MONAILabelApp):
     def __init__(self, app_dir, studies, conf):
         self.models = get_bundle_models(app_dir, conf)
+        # Add Epistemic model for scoring
+        self.epistemic_models = (
+            get_bundle_models(app_dir, conf, conf_key="epistemic_model") if conf.get("epistemic_model") else None
+        )
+        if self.epistemic_models:
+            # Get epistemic parameters
+            self.epistemic_max_samples = int(conf.get("epistemic_max_samples", "0"))
+            self.epistemic_simulation_size = int(conf.get("epistemic_simulation_size", "5"))
+            self.epistemic_dropout = float(conf.get("epistemic_dropout", "0.2"))
 
         super().__init__(
             app_dir=app_dir,
@@ -74,6 +87,32 @@ def init_strategies(self) -> Dict[str, Strategy]:
         logger.info(f"Active Learning Strategies:: {list(strategies.keys())}")
         return strategies
 
+    def init_scoring_methods(self) -> Dict[str, ScoringMethod]:
+        methods: Dict[str, ScoringMethod] = {}
+        if not self.conf.get("epistemic_model"):
+            return methods
+
+        for n, b in self.epistemic_models.items():
+            # Create BundleInferTask task with dropout instantiation for scoring inference
+            i = BundleInferTask(
+                b,
+                self.conf,
+                train_mode=True,
+                skip_writer=True,
+                dropout=self.epistemic_dropout,
+                post_filter=[SaveImaged, Invertd],
+            )
+            methods[n] = EpistemicScoring(
+                i, max_samples=self.epistemic_max_samples, simulation_size=self.epistemic_simulation_size
+            )
+            if not methods:
+                continue
+            methods = methods if isinstance(methods, dict) else {n: methods[n]}
+            logger.info(f"+++ Adding Scoring Method:: {n} => {b}")
+
+        logger.info(f"Active Learning Scoring Methods:: {list(methods.keys())}")
+        return methods
+
 
 """
 Example to run train/infer/scoring task(s) locally without actually running MONAI Label Server