Update infer bundle class to create deepedit infer class

monailabel/tasks/infer/bundle.py

@@ -109,6 +109,8 @@ def __init__(
         self.bundle_path = path
         self.bundle_config_path = os.path.join(path, "configs", config_paths[0])
         self.bundle_config = self._load_bundle_config(self.bundle_path, self.bundle_config_path)
+        # For deepedit inferer - allow the use of clicks
+        self.bundle_config.config["use_click"] = True if type.lower() == "deepedit" else False
 
         if self.dropout > 0:
             self.bundle_config["network_def"]["dropout"] = self.dropout
@@ -133,7 +135,13 @@ def __init__(
         self.key_image, image = next(iter(metadata["network_data_format"]["inputs"].items()))
         self.key_pred, pred = next(iter(metadata["network_data_format"]["outputs"].items()))
 
-        labels = {v.lower(): int(k) for k, v in pred.get("channel_def", {}).items() if v.lower() != "background"}
+        # labels = ({v.lower(): int(k) for k, v in pred.get("channel_def", {}).items() if v.lower() != "background"})
+        labels = {}
+        for k, v in pred.get("channel_def", {}).items():
+            if (not type.lower() == "deepedit") and (v.lower() != "background"):
+                labels[v.lower()] = int(k)
+            else:
+                labels[v.lower()] = int(k)
         description = metadata.get("description")
         spatial_shape = image.get("spatial_shape")
         dimension = len(spatial_shape) if spatial_shape else 3
@@ -192,6 +200,7 @@ def pre_transforms(self, data=None) -> Sequence[Callable]:
             if self.bundle_config.get(k):
                 c = self.bundle_config.get_parsed_content(k, instantiate=True)
                 pre = list(c.transforms) if isinstance(c, Compose) else c
+
         pre = self._filter_transforms(pre, self.pre_filter)
 
         for t in pre:
@@ -254,6 +263,7 @@ def post_transforms(self, data=None) -> Sequence[Callable]:
             if self.bundle_config.get(k):
                 c = self.bundle_config.get_parsed_content(k, instantiate=True)
                 post = list(c.transforms) if isinstance(c, Compose) else c
+
         post = self._filter_transforms(post, self.post_filter)
 
         if self.add_post_restore:

monailabel/utils/others/generic.py

@@ -263,7 +263,7 @@ def get_zoo_bundle(model_dir, conf, models, conf_key):
         print("")
         print("---------------------------------------------------------------------------------------")
         print(
-            "Github access rate limit reached, pleaes provide personal auth token by setting env MONAI_ZOO_AUTH_TOKEN"
+            "Github access rate limit reached, please provide personal auth token by setting env MONAI_ZOO_AUTH_TOKEN"
         )
         print("or --conf auth_token <personal auth token>")
         exit(-1)

sample-apps/monaibundle/main.py

@@ -58,10 +58,22 @@ def init_infers(self) -> Dict[str, InferTask]:
         #################################################
         # Models
         #################################################
+
         for n, b in self.models.items():
-            i = BundleInferTask(b, self.conf)
-            logger.info(f"+++ Adding Inferer:: {n} => {i}")
-            infers[n] = i
+            if "deepedit" in n:
+                # Adding automatic inferer
+                i = BundleInferTask(b, self.conf, type="segmentation")
+                logger.info(f"+++ Adding Inferer:: {n}_seg => {i}")
+                infers[n + "_seg"] = i
+                # Adding inferer for managing clicks
+                i = BundleInferTask(b, self.conf, type="deepedit")
+                logger.info("+++ Adding DeepEdit Inferer")
+                infers[n] = i
+            else:
+                i = BundleInferTask(b, self.conf)
+                logger.info(f"+++ Adding Inferer:: {n} => {i}")
+                infers[n] = i
+
         return infers
 
     def init_trainers(self) -> Dict[str, TrainTask]:
@@ -147,8 +159,9 @@ def main():
     app_dir = os.path.dirname(__file__)
     studies = args.studies
 
-    app = MyApp(app_dir, studies, {"preload": "true", "models": "spleen_ct_segmentation"})
-    train(app)
+    app = MyApp(app_dir, studies, {"preload": "false", "models": "spleen_deepedit_annotation"})
+    # train(app)
+    infer(app)
 
 
 def infer(app):
@@ -157,7 +170,7 @@ def infer(app):
 
     res = app.infer(
         request={
-            "model": "spleen_ct_segmentation",
+            "model": "spleen_deepedit_annotation",
             "image": "image",
         }
     )
@@ -170,7 +183,7 @@ def infer(app):
 def train(app):
     app.train(
         request={
-            "model": "spleen_ct_segmentation",
+            "model": "spleen_deepedit_annotation",
             "max_epochs": 2,
             "multi_gpu": False,
             "val_split": 0.1,