Add Epistemic strategy to DeepEdit App

monailabel/tasks/scoring/epistemic.py

@@ -26,11 +26,11 @@
 
 class EpistemicScoring(ScoringMethod):
     """
-    First version of test time augmentation active learning
+    First version of Epistemic computation used as active learning strategy
     """
 
     def __init__(self, model, network=None, transforms=None, roi_size=(128, 128, 64), num_samples=10):
-        super().__init__("Compute initial score based on TTA")
+        super().__init__("Compute initial score based on dropout")
         self.model = model
         self.network = network
         self.transforms = transforms
@@ -54,7 +54,7 @@ def infer_seg(self, data, model, roi_size, sw_batch_size):
                 inputs=data["image"][None].cuda(), roi_size=roi_size, sw_batch_size=sw_batch_size, predictor=model
             )
 
-        soft_preds = torch.softmax(preds, dim=1)
+        soft_preds = torch.softmax(preds, dim=1) if preds.shape[1] > 1 else torch.sigmoid(preds)
         soft_preds = soft_preds.detach().to("cpu").numpy()
         return soft_preds
 
@@ -125,7 +125,7 @@ def __call__(self, request, datastore: Datastore):
         model, model_ts = self._load_model(self.model, self.network)
         if not model:
             return
-        model = model.to(self.device)
+        model = model.to(self.device).train()
 
         # Performing Epistemic for all unlabeled images
         skipped = 0
@@ -155,7 +155,7 @@ def __call__(self, request, datastore: Datastore):
 
             accum_numpy = np.stack(accum_unl_outputs)
             accum_numpy = np.squeeze(accum_numpy)
-            accum_numpy = accum_numpy[:, 1, :, :, :]
+            accum_numpy = accum_numpy[:, 1:, :, :, :] if len(accum_numpy.shape) > 4 else accum_numpy
 
             entropy = self.entropy_3d_volume(accum_numpy)
             entropy_sum = float(np.sum(entropy))

monailabel/tasks/scoring/tta.py

@@ -141,7 +141,7 @@ def __call__(self, request, datastore: Datastore):
         model, model_ts = self._load_model(self.model, self.network)
         if not model:
             return
-        model = model.to(self.device)
+        model = model.to(self.device).eval()
 
         tt_aug = TestTimeAugmentation(
             transform=self.pre_transforms(),

requirements.txt

@@ -1,6 +1,6 @@
 aiofiles==0.6.0
 fastapi==0.65.2
-monai-weekly[nibabel, skimage, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, lmdb, psutil, openslide]
+monai[nibabel, skimage, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, lmdb, psutil, openslide]==0.7.0rc1
 pyyaml==5.4.1
 python-multipart==0.0.5
 requests-toolbelt==0.9.1

sample-apps/README.md

@@ -12,7 +12,7 @@ the [deepgrow App](./deepgrow). The latter one is meant for the users that want
 
 #### DeepEdit
 
-Similar to the deepgrow Apps, you'll find the one generic [deepedit](./generic_deepedit) that researchers can use to build their own deepedit-based app.
+Similar to the deepgrow Apps, you'll find the one generic [deepedit](./deepedit) that researchers can use to build their own deepedit-based app.
 
 
 #### Automated Segmentation

sample-apps/deepedit/main.py

@@ -25,9 +25,11 @@
 from monailabel.interfaces.tasks.strategy import Strategy
 from monailabel.interfaces.tasks.train import TrainTask
 from monailabel.scribbles.infer import HistogramBasedGraphCut
+from monailabel.tasks.activelearning.epistemic import Epistemic
 from monailabel.tasks.activelearning.random import Random
 from monailabel.tasks.activelearning.tta import TTA
 from monailabel.tasks.scoring.dice import Dice
+from monailabel.tasks.scoring.epistemic import EpistemicScoring
 from monailabel.tasks.scoring.sum import Sum
 from monailabel.tasks.scoring.tta import TTAScoring
 from monailabel.utils.others.planner import HeuristicPlanner
@@ -37,37 +39,39 @@
 
 class MyApp(MONAILabelApp):
     def __init__(self, app_dir, studies, conf):
-        self.network = DynUNetV1(
-            spatial_dims=3,
-            in_channels=3,
-            out_channels=1,
-            kernel_size=[
+        network_params = {
+            "spatial_dims": 3,
+            "in_channels": 3,
+            "out_channels": 1,
+            "kernel_size": [
                 [3, 3, 3],
                 [3, 3, 3],
                 [3, 3, 3],
                 [3, 3, 3],
                 [3, 3, 3],
                 [3, 3, 3],
             ],
-            strides=[
+            "strides": [
                 [1, 1, 1],
                 [2, 2, 2],
                 [2, 2, 2],
                 [2, 2, 2],
                 [2, 2, 2],
                 [2, 2, 1],
             ],
-            upsample_kernel_size=[
+            "upsample_kernel_size": [
                 [2, 2, 2],
                 [2, 2, 2],
                 [2, 2, 2],
                 [2, 2, 2],
                 [2, 2, 1],
             ],
-            norm_name="instance",
-            deep_supervision=False,
-            res_block=True,
-        )
+            "norm_name": "instance",
+            "deep_supervision": False,
+            "res_block": True,
+        }
+        self.network = DynUNetV1(**network_params)
+        self.network_with_dropout = DynUNetV1(**network_params, dropout=0.2)
 
         self.model_dir = os.path.join(app_dir, "model")
         self.pretrained_model = os.path.join(self.model_dir, "pretrained.pt")
@@ -86,6 +90,10 @@ def __init__(self, app_dir, studies, conf):
         if use_pretrained_model:
             self.download([(self.pretrained_model, pretrained_model_uri)])
 
+        self.epistemic_enabled = strtobool(conf.get("epistemic_enabled", "false"))
+        self.epistemic_samples = int(conf.get("epistemic_samples", "5"))
+        logger.info(f"EPISTEMIC Enabled: {self.epistemic_enabled}; Samples: {self.epistemic_samples}")
+
         self.tta_enabled = strtobool(conf.get("tta_enabled", "false"))
         self.tta_samples = int(conf.get("tta_samples", "5"))
         logger.info(f"TTA Enabled: {self.tta_enabled}; Samples: {self.tta_samples}")
@@ -137,6 +145,8 @@ def init_trainers(self) -> Dict[str, TrainTask]:
 
     def init_strategies(self) -> Dict[str, Strategy]:
         strategies: Dict[str, Strategy] = {}
+        if self.epistemic_enabled:
+            strategies["EPISTEMIC"] = Epistemic()
         if self.tta_enabled:
             strategies["TTA"] = TTA()
         strategies["random"] = Random()
@@ -145,6 +155,13 @@ def init_strategies(self) -> Dict[str, Strategy]:
 
     def init_scoring_methods(self) -> Dict[str, ScoringMethod]:
         methods: Dict[str, ScoringMethod] = {}
+        if self.epistemic_enabled:
+            methods["EPISTEMIC"] = EpistemicScoring(
+                model=[self.pretrained_model, self.final_model],
+                network=self.network_with_dropout,
+                transforms=self._infers["deepedit_seg"].pre_transforms(),
+                num_samples=self.epistemic_samples,
+            )
         if self.tta_enabled:
             methods["TTA"] = TTAScoring(
                 model=[self.pretrained_model, self.final_model],