diff --git a/monailabel/tasks/scoring/epistemic.py b/monailabel/tasks/scoring/epistemic.py index 6bea0ace5..31a942cdd 100644 --- a/monailabel/tasks/scoring/epistemic.py +++ b/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)) diff --git a/monailabel/tasks/scoring/tta.py b/monailabel/tasks/scoring/tta.py index d2ef1a468..604255ddc 100644 --- a/monailabel/tasks/scoring/tta.py +++ b/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(), diff --git a/requirements.txt b/requirements.txt index 19c814ddd..940e912d4 100644 --- a/requirements.txt +++ b/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 diff --git a/sample-apps/README.md b/sample-apps/README.md index 3ab5e42a6..92fb6e37f 100644 --- a/sample-apps/README.md +++ b/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 diff --git a/sample-apps/deepedit/main.py b/sample-apps/deepedit/main.py index 7d3823194..04543a8e7 100644 --- a/sample-apps/deepedit/main.py +++ b/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,11 +39,11 @@ 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], @@ -49,7 +51,7 @@ def __init__(self, app_dir, studies, conf): [3, 3, 3], [3, 3, 3], ], - strides=[ + "strides": [ [1, 1, 1], [2, 2, 2], [2, 2, 2], @@ -57,17 +59,19 @@ def __init__(self, app_dir, studies, conf): [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],