UI overhaul + overlap parameter

- Added overlap parameter for window
- Improved UI code slightly

Author: C-Achard

napari_cellseg3d/interface.py

@@ -1,5 +1,4 @@
 from typing import Optional
-from typing import Union
 from typing import List
 
 
@@ -61,6 +60,13 @@ def toggle_visibility(checkbox, widget):
     widget.setVisible(checkbox.isChecked())
 
 
+def add_label(widget, label, label_before=True, horizontal=True):
+    if label_before:
+        return combine_blocks(widget, label, horizontal=horizontal)
+    else:
+        return combine_blocks(label, widget, horizontal=horizontal)
+
+
 class Button(QPushButton):
     """Class for a button with a title and connected to a function when clicked. Inherits from QPushButton.
 
@@ -494,20 +500,33 @@ def __init__(
         step=1,
         parent: Optional[QWidget] = None,
         fixed: Optional[bool] = True,
+        label: Optional[str] = None,
     ):
         """Args:
         min (Optional[float]): minimum value, defaults to 0
         max (Optional[float]): maximum value, defaults to 10
         default (Optional[float]): default value, defaults to 0
         step (Optional[float]): step value, defaults to 1
         parent: parent widget, defaults to None
-        fixed (bool): if True, sets the QSizePolicy of the spinbox to Fixed"""
+        fixed (bool): if True, sets the QSizePolicy of the spinbox to Fixed
+        label (Optional[str]): if provided, creates a label with the chosen title to use with the counter"""
 
         super().__init__(parent)
         set_spinbox(self, min, max, default, step, fixed)
 
+        if label is not None:
+            self.label = make_label(name=label)
+
+    # def setToolTip(self, a0: str) -> None:
+    #     self.setToolTip(a0)
+    #     if self.label is not None:
+    #         self.label.setToolTip(a0)
+
+    def get_with_label(self, horizontal=True):
+        return add_label(self, self.label, horizontal=horizontal)
+
     def set_precision(self, decimals):
-        """Sets the precision of the box to the speicifed number of decimals"""
+        """Sets the precision of the box to the specified number of decimals"""
         self.setDecimals(decimals)
 
     @classmethod
@@ -535,6 +554,7 @@ def __init__(
         step=1,
         parent: Optional[QWidget] = None,
         fixed: Optional[bool] = True,
+        label: Optional[str] = None,
     ):
         """Args:
         min (Optional[int]): minimum value, defaults to 0
@@ -546,6 +566,9 @@ def __init__(
 
         super().__init__(parent)
         set_spinbox(self, min, max, default, step, fixed)
+        self.label = None
+        if label is not None:
+            self.label = make_label(label, self)
 
     @classmethod
     def make_n(

napari_cellseg3d/model_framework.py

@@ -14,6 +14,7 @@
 from napari_cellseg3d.log_utility import Log
 from napari_cellseg3d.models import model_SegResNet as SegResNet
 from napari_cellseg3d.models import model_SwinUNetR as SwinUNetR
+
 # from napari_cellseg3d.models import model_TRAILMAP as TRAILMAP
 from napari_cellseg3d.models import model_VNet as VNet
 from napari_cellseg3d.models import model_TRAILMAP_MS as TRAILMAP_MS

napari_cellseg3d/model_workers.py

@@ -164,6 +164,9 @@ def __init__(self):
         super().__init__()
 
 
+# TODO : use dataclass for config instead ?
+
+
 class InferenceWorker(GeneratorWorker):
     """A custom worker to run inference jobs in.
     Inherits from :py:class:`napari.qt.threading.GeneratorWorker`"""
@@ -179,6 +182,7 @@ def __init__(
         instance,
         use_window,
         window_infer_size,
+        window_overlap,
         keep_on_cpu,
         stats_csv,
         images_filepaths=None,
@@ -230,7 +234,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.window_overlap_percentage = window_overlap
         self.keep_on_cpu = keep_on_cpu
         self.stats_to_csv = stats_csv
         ############################################
@@ -315,17 +319,53 @@ def load_folder(self):
         self.log("\nChecking dimensions...")
         pad = utils.get_padding_dim(check)
 
-        load_transforms = Compose(
-            [
-                LoadImaged(keys=["image"]),
-                # AddChanneld(keys=["image"]), #already done
-                EnsureChannelFirstd(keys=["image"]),
-                # Orientationd(keys=["image"], axcodes="PLI"),
-                # anisotropic_transform,
-                SpatialPadd(keys=["image"], spatial_size=pad),
-                EnsureTyped(keys=["image"]),
-            ]
-        )
+        dims = self.model_dict["model_input_size"]
+
+        if self.model_dict["name"] == "SegResNet":
+            model = self.model_dict["class"].get_net(
+                input_image_size=[
+                    dims,
+                    dims,
+                    dims,
+                ]
+            )
+        elif self.model_dict["name"] == "SwinUNetR":
+            model = self.model_dict["class"].get_net(
+                img_size=[dims, dims, dims],
+                use_checkpoint=False,
+            )
+        else:
+            model = self.model_dict["class"].get_net()
+
+        self.log_parameters()
+
+        model.to(self.device)
+
+        # print("FILEPATHS PRINT")
+        # print(self.images_filepaths)
+        if self.use_window:
+            load_transforms = Compose(
+                [
+                    LoadImaged(keys=["image"]),
+                    # AddChanneld(keys=["image"]), #already done
+                    EnsureChannelFirstd(keys=["image"]),
+                    # Orientationd(keys=["image"], axcodes="PLI"),
+                    # anisotropic_transform,
+                    EnsureTyped(keys=["image"]),
+                ]
+            )
+        else:
+            load_transforms = Compose(
+                [
+                    LoadImaged(keys=["image"]),
+                    # AddChanneld(keys=["image"]), #already done
+                    EnsureChannelFirstd(keys=["image"]),
+                    # Orientationd(keys=["image"], axcodes="PLI"),
+                    # anisotropic_transform,
+                    SpatialPadd(keys=["image"], spatial_size=pad),
+                    EnsureTyped(keys=["image"]),
+                ]
+            )
 
         self.log("\nLoading dataset...")
         inference_ds = Dataset(data=images_dict, transform=load_transforms)

napari_cellseg3d/plugin_model_inference.py

@@ -78,7 +78,7 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
         self.keep_on_cpu = False
         self.use_window_inference = False
         self.window_inference_size = None
-        self.window_overlap_percentage = None
+        self.window_overlap = 0.25
 
         ###########################
         # interface
@@ -131,7 +131,7 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
             T=7, parent=self
         )
 
-        self.window_infer_box = ui.make_checkbox("Use window inference")
+        self.window_infer_box = ui.CheckBox(title="Use window inference")
         self.window_infer_box.clicked.connect(self.toggle_display_window_size)
 
         sizes_window = ["8", "16", "32", "64", "128", "256", "512"]
@@ -151,9 +151,18 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
         )
         self.lbl_window_size_choice = self.window_size_choice.label
 
-        self.keep_data_on_cpu_box = ui.make_checkbox("Keep data on CPU")
+        self.window_overlap_counter = ui.DoubleIncrementCounter(
+            min=0,
+            max=1,
+            default=0.25,
+            step=0.05,
+            parent=self,
+            label="Overlap %",
+        )
 
-        self.window_infer_params = ui.combine_blocks(
+        self.keep_data_on_cpu_box = ui.CheckBox(title="Keep data on CPU")
+
+        window_size_widgets = ui.combine_blocks(
             self.window_size_choice,
             self.lbl_window_size_choice,
             horizontal=False,
@@ -164,6 +173,12 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
         #     horizontal=False,
         # )
 
+        self.window_infer_params = ui.combine_blocks(
+            window_size_widgets,
+            self.window_overlap_counter.get_with_label(horizontal=False),
+            horizontal=False,
+        )
+
         ##################
         ##################
         # instance segmentation widgets
@@ -252,6 +267,10 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
             "Size of the window to run inference with (in pixels)"
         )
 
+        self.window_overlap_counter.setToolTip(
+            "Percentage of overlap between windows to use when using sliding window"
+        )
+
         # self.window_overlap.setToolTip(
         #     "Amount of overlap between sliding windows"
         # )
@@ -625,7 +644,7 @@ def start(self, on_layer=False):
             self.window_inference_size = int(
                 self.window_size_choice.currentText()
             )
-            # self.window_overlap_percentage = self.window_overlap.value()
+            self.window_overlap = self.window_overlap_counter.value()
 
             if not on_layer:
                 self.worker = InferenceWorker(
@@ -639,6 +658,7 @@ def start(self, on_layer=False):
                     instance=self.instance_params,
                     use_window=self.use_window_inference,
                     window_infer_size=self.window_inference_size,
+                    window_overlap=self.window_overlap,
                     keep_on_cpu=self.keep_on_cpu,
                     stats_csv=self.stats_to_csv,
                 )
@@ -655,6 +675,7 @@ def start(self, on_layer=False):
                     use_window=self.use_window_inference,
                     window_infer_size=self.window_inference_size,
                     keep_on_cpu=self.keep_on_cpu,
+                    window_overlap=self.window_overlap,
                     stats_csv=self.stats_to_csv,
                     layer=layer,
                 )