adding swinunetr & removing extra padding for inference

Author: vidalmaxime

.gitignore

@@ -99,4 +99,3 @@ venv/
 /napari_cellseg3d/models/saved_weights/
 /docs/res/logo/old_logo/
 /reqs/
-

napari_cellseg3d/model_framework.py

@@ -13,6 +13,7 @@
 from napari_cellseg3d import utils
 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
@@ -64,6 +65,7 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
             "SegResNet": SegResNet,
             "TRAILMAP": TRAILMAP,
             "TRAILMAP_MS": TRAILMAP_MS,
+            "SwinUNetR": SwinUNetR,
         }
         """dict: dictionary of available models, with string for widget display as key
 

napari_cellseg3d/model_workers.py

@@ -301,8 +301,6 @@ def log_parameters(self):
                 f"Probability threshold is {self.instance_params['threshold']:.2f}\n"
                 f"Objects smaller than {self.instance_params['size_small']} pixels will be removed\n"
             )
-            # self.log(f"")
-            # self.log("\n")
         self.log("-" * 20)
 
     def load_folder(self):
@@ -313,11 +311,7 @@ def load_folder(self):
         data_check = LoadImaged(keys=["image"])(images_dict[0])
 
         check = data_check["image"].shape
-        # TODO remove
-        # z_aniso = 5 / 1.5
-        # if zoom is not None :
-        #     pad = utils.get_padding_dim(check, anisotropy_factor=zoom)
-        # else:
+
         self.log("\nChecking dimensions...")
         pad = utils.get_padding_dim(check)
 
@@ -1027,10 +1021,26 @@ def train(self):
                     out_channels=1,
                     dropout_prob=0.3,
                 )
+            elif model_name == "SwinUNetR":
+                if self.sampling:
+                    size = self.sample_size
+                else:
+                    size = check
+                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:
                 model = model_class.get_net()  # get an instance of the model
             model = model.to(self.device)
 
+
+
+
             epoch_loss_values = []
             val_metric_values = []
 

napari_cellseg3d/models/model_SwinUNetR.py

@@ -0,0 +1,18 @@
+from monai.networks.nets import SwinUNETR
+
+
+def get_weights_file():
+    return ""
+
+
+def get_net():
+    return SwinUNETR
+
+
+def get_output(model, input):
+    out = model(input)
+    return out
+
+
+def get_validation(model, val_inputs):
+    return model(val_inputs)

napari_cellseg3d/models/model_VNet.py

@@ -19,6 +19,11 @@ def get_validation(model, val_inputs):
     roi_size = (64, 64, 64)
     sw_batch_size = 1
     val_outputs = sliding_window_inference(
-        val_inputs, roi_size, sw_batch_size, model, mode="gaussian"
+        val_inputs,
+        roi_size,
+        sw_batch_size,
+        model,
+        mode="gaussian",
+        overlap=0.7,
     )
     return val_outputs

napari_cellseg3d/plugin_model_inference.py

@@ -78,6 +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
 
         ###########################
         # interface
@@ -134,6 +135,17 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
         self.window_infer_box.clicked.connect(self.toggle_display_window_size)
 
         sizes_window = ["8", "16", "32", "64", "128", "256", "512"]
+        # (
+        #     self.window_size_choice,
+        #     self.lbl_window_size_choice,
+        # ) = ui.make_combobox(sizes_window, label="Window size and overlap")
+        # self.window_overlap = ui.make_n_spinboxes(
+        #     max=1,
+        #     default=0.7,
+        #     step=0.05,
+        #     double=True,
+        # )
+
         self.window_size_choice = ui.DropdownMenu(
             sizes_window, label="Window size"
         )
@@ -146,6 +158,11 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
             self.lbl_window_size_choice,
             horizontal=False,
         )
+        # self.window_infer_params = ui.combine_blocks(
+        #     self.window_overlap,
+        #     self.window_infer_params,
+        #     horizontal=False,
+        # )
 
         ##################
         ##################
@@ -216,7 +233,7 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
             "Displays the image used for inference in the viewer"
         )
         self.segres_size.setToolTip(
-            "Image size on which the SegResNet has been trained (default : 128)"
+            "Image size on which the model has been trained (default : 128)"
         )
 
         thresh_desc = (
@@ -234,6 +251,11 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
         self.window_size_choice.setToolTip(
             "Size of the window to run inference with (in pixels)"
         )
+
+        # self.window_overlap.setToolTip(
+        #     "Amount of overlap between sliding windows"
+        # )
+
         self.keep_data_on_cpu_box.setToolTip(
             "If enabled, data will be kept on the RAM rather than the VRAM.\nCan avoid out of memory issues with CUDA"
         )
@@ -281,7 +303,10 @@ def check_ready(self):
             return False
 
     def toggle_display_segres_size(self):
-        if self.model_choice.currentText() == "SegResNet":
+        if (
+            self.model_choice.currentText() == "SegResNet"
+            or self.model_choice.currentText() == "SwinUNetR"
+        ):
             self.segres_size.setVisible(True)
         else:
             self.segres_size.setVisible(False)
@@ -600,6 +625,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()
 
             if not on_layer:
                 self.worker = InferenceWorker(
@@ -724,8 +750,6 @@ def on_yield(data, widget):
 
             zoom = widget.zoom
 
-            # print(data["original"].shape)
-            # print(data["result"].shape)
 
             viewer.dims.ndisplay = 3
             viewer.scale_bar.visible = True

napari_cellseg3d/plugin_model_training.py

@@ -31,7 +31,7 @@
 from napari_cellseg3d.model_workers import TrainingWorker
 
 NUMBER_TABS = 3
-DEFAULT_PATCH_SIZE = 60
+DEFAULT_PATCH_SIZE = 64
 
 
 class Trainer(ModelFramework):

napari_cellseg3d/utils.py

@@ -116,7 +116,6 @@ def resize(image, zoom_factors):
 
 
 def align_array_sizes(array_shape, target_shape):
-
     index_differences = []
     for i in range(len(target_shape)):
         if target_shape[i] != array_shape[i]:
@@ -334,7 +333,6 @@ def fill_list_in_between(lst, n, elem):
     Returns :
         Filled list
     """
-
     new_list = []
     for i in range(len(lst)):
         temp_list = [lst[i]]
@@ -608,25 +606,6 @@ def format_Warning(message, category, filename, lineno, line=""):
     )
 
 
-# def dice_coeff(y_true, y_pred):
-#     smooth = 1.
-#     y_true_f = y_true.flatten()
-#     y_pred_f = K.flatten(y_pred)
-#     intersection = K.sum(y_true_f * y_pred_f)
-#     score = (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)
-#     return score
-
-
-# def dice_loss(y_true, y_pred):
-#     loss = 1 - dice_coeff(y_true, y_pred)
-#     return loss
-
-
-# def bce_dice_loss(y_true, y_pred):
-#     loss = binary_crossentropy(y_true, y_pred) + dice_loss(y_true, y_pred)
-#     return loss
-
-
 def divide_imgs(images):
     H = -(-images.shape[1] // 412)
     W = -(-images.shape[2] // 412)

setup.cfg

@@ -50,12 +50,13 @@ install_requires =
     tifffile>=2022.2.9
     imageio-ffmpeg>=0.4.5
     torch>=1.11
+    monai[nibabel,scikit-image,itk,einops]>=0.9.0
     tqdm
     monai>=0.9.0
     nibabel
     scikit-image
     pillow
-    itk>=5.2.0
+    tqdm
     matplotlib
     vispy>=0.9.6