Change threshold for Swin

Author: C-Achard

napari_cellseg3d/code_models/worker_inference.py

@@ -9,7 +9,7 @@
 from monai.inferers import sliding_window_inference
 from monai.transforms import (
     AddChannel,
-    AsDiscrete,
+    # AsDiscrete,
     Compose,
     EnsureChannelFirstd,
     EnsureType,
@@ -35,6 +35,7 @@
     QuantileNormalization,
     QuantileNormalizationd,
     RemapTensor,
+    Threshold,
     WeightsDownloader,
 )
 
@@ -727,7 +728,8 @@ def inference(self):
                 post_process_transforms = Compose(
                     [
                         RemapTensor(new_max=1.0, new_min=0.0),
-                        AsDiscrete(threshold=t),
+                        # AsDiscrete(threshold=t),
+                        Threshold(threshold=t),
                         EnsureType(),
                     ]
                 )

napari_cellseg3d/code_models/worker_training.py

@@ -17,7 +17,7 @@
 from monai.inferers import sliding_window_inference
 from monai.metrics import DiceMetric
 from monai.transforms import (
-    AsDiscrete,
+    # AsDiscrete,
     Compose,
     EnsureChannelFirstd,
     EnsureType,
@@ -45,6 +45,7 @@
     LogSignal,
     QuantileNormalizationd,
     RemapTensor,
+    Threshold,
     TrainingReport,
     WeightsDownloader,
 )
@@ -638,10 +639,11 @@ def get_loader_func(num_samples):
                             labs = decollate_batch(val_labels)
 
                             # TODO : more parameters/flexibility
+
                             post_pred = Compose(
                                 [
                                     RemapTensor(new_max=1, new_min=0),
-                                    AsDiscrete(threshold=0.25),  # needed ?
+                                    Threshold(threshold=0.5),
                                     EnsureType(),
                                 ]
                             )  #

napari_cellseg3d/code_models/workers_utils.py

@@ -213,6 +213,15 @@ def __call__(self, img):
         return utils.remap_image(img, new_max=self.max, new_min=self.min)
 
 
+class Threshold(Transform):
+    def __init__(self, threshold=0.5):
+        super().__init__()
+        self.threshold = threshold
+
+    def __call__(self, img):
+        return torch.where(img > self.threshold, 1, 0)
+
+
 @dataclass
 class InferenceResult:
     """Class to record results of a segmentation job"""