:wrench move inherent model params where they are instantiated and fixed padding only for sliding window inference

Author: vidalmaxime

napari_cellseg3d/model_workers.py

@@ -45,7 +45,6 @@
 # Qt
 from qtpy.QtCore import Signal
 
-
 from napari_cellseg3d import utils
 from napari_cellseg3d import log_utility
 
@@ -231,6 +230,7 @@ def __init__(
         self.instance_params = instance
         self.use_window = use_window
         self.window_infer_size = window_infer_size
+        self.window_overlap_percentage = 0.8,
         self.keep_on_cpu = keep_on_cpu
         self.stats_to_csv = stats_csv
         ############################################
@@ -399,8 +399,10 @@ def model_output(
 
         if self.use_window:
             window_size = self.window_infer_size
+            window_overlap = self.window_overlap_percentage
         else:
             window_size = None
+            window_overlap = 0.25
 
         outputs = sliding_window_inference(
             inputs,
@@ -409,6 +411,7 @@ def model_output(
             predictor=model_output,
             sw_device=self.device,
             device=dataset_device,
+            overlap=window_overlap,
         )
 
         out = outputs.detach().cpu()
@@ -1029,9 +1032,6 @@ def train(self):
                 print(f"Size of image : {size}")
                 model = model_class.get_net()(
                     img_size=utils.get_padding_dim(size),
-                    in_channels=1,
-                    out_channels=1,
-                    feature_size=48,
                     use_checkpoint=True,
                 )
             else:

napari_cellseg3d/models/model_SegResNet.py

@@ -1,8 +1,8 @@
 from monai.networks.nets import SegResNetVAE
 
 
-def get_net():
-    return SegResNetVAE
+def get_net(input_image_size, dropout_prob=None):
+    return SegResNetVAE(input_image_size, out_channels=1, dropout_prob=dropout_prob)
 
 
 def get_weights_file():

napari_cellseg3d/models/model_SwinUNetR.py

@@ -1,17 +1,18 @@
+import torch
 from monai.networks.nets import SwinUNETR
 
 
 def get_weights_file():
     return ""
 
 
-def get_net():
-    return SwinUNETR
+def get_net(img_size, use_checkpoint=True):
+    return SwinUNETR(img_size, in_channels=1, out_channels=1, feature_size=48, use_checkpoint=use_checkpoint)
 
 
 def get_output(model, input):
     out = model(input)
-    return out
+    return torch.sigmoid(out)
 
 
 def get_validation(model, val_inputs):

napari_cellseg3d/plugin_model_inference.py

@@ -100,9 +100,9 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
         ######################
         ######################
         # TODO : better way to handle SegResNet size reqs ?
-        self.segres_size = ui.IntIncrementCounter(min=1, max=1024, default=128)
+        self.model_input_size = ui.IntIncrementCounter(min=1, max=1024, default=128)
         self.model_choice.currentIndexChanged.connect(
-            self.toggle_display_segres_size
+            self.toggle_display_model_input_size
         )
         self.model_choice.setCurrentIndex(0)
 
@@ -232,7 +232,7 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
         self.show_original_checkbox.setToolTip(
             "Displays the image used for inference in the viewer"
         )
-        self.segres_size.setToolTip(
+        self.model_input_size.setToolTip(
             "Image size on which the model has been trained (default : 128)"
         )
 
@@ -302,14 +302,14 @@ def check_ready(self):
             warnings.warn("Image and label paths are not correctly set")
             return False
 
-    def toggle_display_segres_size(self):
+    def toggle_display_model_input_size(self):
         if (
             self.model_choice.currentText() == "SegResNet"
             or self.model_choice.currentText() == "SwinUNetR"
         ):
-            self.segres_size.setVisible(True)
+            self.model_input_size.setVisible(True)
         else:
-            self.segres_size.setVisible(False)
+            self.model_input_size.setVisible(False)
 
     def toggle_display_number(self):
         """Shows the choices for viewing results depending on whether :py:attr:`self.view_checkbox` is checked"""
@@ -418,7 +418,7 @@ def build(self):
                 self.model_choice,
                 self.custom_weights_choice,
                 self.weights_path_container,
-                self.segres_size,
+                self.model_input_size,
             ],
         )
         self.weights_path_container.setVisible(False)
@@ -576,7 +576,7 @@ def start(self, on_layer=False):
             model_dict = {  # gather model info
                 "name": model_key,
                 "class": self.get_model(model_key),
-                "segres_size": self.segres_size.value(),
+                "model_input_size": self.model_input_size.value(),
             }
 
             if self.custom_weights_choice.isChecked():