Enable window + default thresholds

Author: C-Achard

napari_cellseg3d/code_models/models/TEMPLATE_model.py

@@ -2,6 +2,7 @@
 
 Please note that custom model implementations are not fully supported out of the box yet, but might be in the future.
 """
+
 from abc import ABC, abstractmethod
 
 
@@ -11,6 +12,7 @@ class ModelTemplate_(ABC):
     weights_file = (
         "model_template.pth"  # specify the file name of the weights file only
     )
+    default_threshold = 0.5  # specify the default threshold for the model
 
     @abstractmethod
     def __init__(

napari_cellseg3d/code_models/models/model_SegResNet.py

@@ -1,11 +1,13 @@
 """SegResNet wrapper for napari_cellseg3d."""
+
 from monai.networks.nets import SegResNetVAE
 
 
 class SegResNet_(SegResNetVAE):
     """SegResNet_ wrapper for napari_cellseg3d."""
 
     weights_file = "SegResNet_latest.pth"
+    default_threshold = 0.3
 
     def __init__(
         self, input_img_size, out_channels=1, dropout_prob=0.3, **kwargs

napari_cellseg3d/code_models/models/model_SwinUNetR.py

@@ -1,4 +1,5 @@
 """SwinUNetR wrapper for napari_cellseg3d."""
+
 from monai.networks.nets import SwinUNETR
 
 from napari_cellseg3d.utils import LOGGER
@@ -10,6 +11,7 @@ class SwinUNETR_(SwinUNETR):
     """SwinUNETR wrapper for napari_cellseg3d."""
 
     weights_file = "SwinUNetR_latest.pth"
+    default_threshold = 0.4
 
     def __init__(
         self,

napari_cellseg3d/code_models/models/model_TRAILMAP_MS.py

@@ -1,4 +1,5 @@
 """TRAILMAP model, reimplemented in PyTorch."""
+
 from napari_cellseg3d.code_models.models.unet.model import UNet3D
 from napari_cellseg3d.utils import LOGGER as logger
 
@@ -7,6 +8,7 @@ class TRAILMAP_MS_(UNet3D):
     """TRAILMAP_MS wrapper for napari_cellseg3d."""
 
     weights_file = "TRAILMAP_MS_best_metric.pth"
+    default_threshold = 0.15
 
     # original model from Liqun Luo lab, transferred to pytorch and trained on mesoSPIM-acquired data (mostly TPH2 as of July 2022)
 

napari_cellseg3d/code_models/models/model_VNet.py

@@ -1,11 +1,13 @@
 """VNet wrapper for napari_cellseg3d."""
+
 from monai.networks.nets import VNet
 
 
 class VNet_(VNet):
     """VNet wrapper for napari_cellseg3d."""
 
     weights_file = "VNet_latest.pth"
+    default_threshold = 0.15
 
     def __init__(self, in_channels=1, out_channels=1, **kwargs):
         """Create a VNet model.

napari_cellseg3d/code_models/models/model_WNet.py

@@ -15,6 +15,7 @@ class WNet_(WNet_encoder):
     """
 
     weights_file = "wnet_latest.pth"
+    default_threshold = 0.6
 
     def __init__(
         self,

napari_cellseg3d/code_models/models/model_test.py

@@ -1,4 +1,5 @@
 """Model for testing purposes."""
+
 import torch
 from torch import nn
 
@@ -7,6 +8,7 @@ class TestModel(nn.Module):
     """For tests only."""
 
     weights_file = "test.pth"
+    default_threshold = 0.5
 
     def __init__(self, **kwargs):
         """Create a TestModel model."""

napari_cellseg3d/code_plugins/plugin_model_inference.py

@@ -143,11 +143,11 @@ def __init__(self, viewer: "napari.viewer.Viewer", parent=None):
         self.thresholding_checkbox = ui.CheckBox(
             "Perform thresholding", self._toggle_display_thresh
         )
-        self.thresholding_checkbox.setChecked(True)
 
         self.thresholding_slider = ui.Slider(
-            default=config.PostProcessConfig().thresholding.threshold_value
-            * 100,
+            default=config.MODEL_LIST[
+                self.model_choice.currentText()
+            ].default_threshold,
             divide_factor=100.0,
             parent=self,
         )
@@ -410,6 +410,13 @@ def _load_weights_path(self):
         )
         self._update_weights_path(file)
 
+    def _set_default_threshold(self):
+        # Whenever a model is selected, set the default threshold from the model file
+        model_name = self.model_choice.currentText()
+        threshold = config.MODEL_LIST[model_name].default_threshold
+        print(threshold)
+        self.thresholding_slider.slider_value = threshold * 100
+
     def _build(self):
         """Puts all widgets in a layout and adds them to the napari Viewer."""
         # ui.add_blank(self.view_results_container, view_results_layout)
@@ -494,7 +501,8 @@ def _build(self):
                 self.device_choice,
             ],
         )
-        self.window_infer_params.setVisible(False)
+        self.use_window_choice.setChecked(True)
+        # self.window_infer_params.setVisible(False)
 
         inference_param_group_w.setLayout(inference_param_group_l)
 
@@ -539,14 +547,18 @@ def _build(self):
                 # self.instance_param_container,  # instance segmentation
             ],
         )
+        # self.thresholding_slider.container.setVisible(False)
+        self.thresholding_checkbox.setChecked(True)
         self._toggle_crf_choice()
         self.model_choice.currentIndexChanged.connect(self._toggle_crf_choice)
+        self.model_choice.currentIndexChanged.connect(
+            self._set_default_threshold
+        )
         ModelFramework._show_io_element(
             self.save_stats_to_csv_box, self.use_instance_choice
         )
 
         self.anisotropy_wdgt.container.setVisible(False)
-        self.thresholding_slider.container.setVisible(False)
         self.instance_widgets.setVisible(False)
         self.crf_widgets.setVisible(False)
         self.save_stats_to_csv_box.setVisible(False)