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

@@ -179,6 +179,7 @@ def __init__(
         instance,
         use_window,
         window_infer_size,
+        window_overlap_percentage,
         keep_on_cpu,
         stats_csv,
     ):
@@ -205,6 +206,8 @@ def __init__(
 
             * window_infer_size: size of window if use_window is True
 
+            * window_overlap_percentage: overlap of sliding windows if use_window is True
+
             * keep_on_cpu: keep images on CPU or no
 
             * stats_csv: compute stats on cells and save them to a csv file
@@ -228,6 +231,7 @@ def __init__(
         self.instance_params = instance
         self.use_window = use_window
         self.window_infer_size = window_infer_size
+        self.window_overlap_percentage = window_overlap_percentage
         self.keep_on_cpu = keep_on_cpu
         self.stats_to_csv = stats_csv
         """These attributes are all arguments of :py:func:~inference, please see that for reference"""
@@ -350,8 +354,6 @@ def inference(self):
         #     pad = utils.get_padding_dim(check, anisotropy_factor=zoom)
         # else:
         self.log("\nChecking dimensions...")
-        pad = utils.get_padding_dim(check)
-        # print(pad)
         dims = self.model_dict["segres_size"]
 
         model = self.model_dict["class"].get_net()
@@ -365,6 +367,14 @@ def inference(self):
                 out_channels=1,
                 # dropout_prob=0.3,
             )
+        elif self.model_dict["name"] == "SwinUNetR":
+            model = self.model_dict["class"].get_net()(
+                img_size=[dims, dims, dims],
+                in_channels=1,
+                out_channels=1,
+                feature_size=48,
+                use_checkpoint=False,
+            )
 
         self.log_parameters()
 
@@ -380,7 +390,6 @@ def inference(self):
                 EnsureChannelFirstd(keys=["image"]),
                 # Orientationd(keys=["image"], axcodes="PLI"),
                 # anisotropic_transform,
-                SpatialPadd(keys=["image"], spatial_size=pad),
                 EnsureTyped(keys=["image"]),
             ]
         )
@@ -437,10 +446,18 @@ def inference(self):
                 # print(inputs.shape)
 
                 inputs = inputs.to("cpu")
+                print(inputs.shape)
 
-                model_output = lambda inputs: post_process_transforms(
-                    self.model_dict["class"].get_output(model, inputs)
-                )
+                if self.model_dict["name"] == "SwinUNetR":
+                    model_output = lambda inputs: post_process_transforms(
+                        torch.sigmoid(
+                            self.model_dict["class"].get_output(model, inputs)
+                        )
+                    )
+                else:
+                    model_output = lambda inputs: post_process_transforms(
+                        self.model_dict["class"].get_output(model, inputs)
+                    )
 
                 if self.keep_on_cpu:
                     dataset_device = "cpu"
@@ -449,22 +466,24 @@ def inference(self):
 
                 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,
                     roi_size=window_size,
                     sw_batch_size=1,
                     predictor=model_output,
                     sw_device=self.device,
                     device=dataset_device,
+                    overlap=window_overlap,
                 )
-
+                print("done window infernce")
                 out = outputs.detach().cpu()
                 # del outputs # TODO fix memory ?
                 # outputs = None
-
+                print(out.shape)
                 if self.transforms["zoom"][0]:
                     zoom = self.transforms["zoom"][1]
                     anisotropic_transform = Zoom(
@@ -474,9 +493,11 @@ def inference(self):
                     )
                     out = anisotropic_transform(out[0])
 
-                out = post_process_transforms(out)
+                # out = post_process_transforms(out)
                 out = np.array(out).astype(np.float32)
+                print(out.shape)
                 out = np.squeeze(out)
+                print(out.shape)
                 to_instance = out  # avoid post processing since thresholding is done there anyway
 
                 # batch_len = out.shape[1]
@@ -825,6 +846,19 @@ 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)

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

@@ -76,6 +76,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
@@ -132,6 +133,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"
         )
@@ -144,6 +156,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,
+        # )
 
         ##################
         ##################
@@ -210,7 +227,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 = "Thresholding : all values in the image below the chosen probability threshold will be set to 0, and all others to 1."
@@ -224,6 +241,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"
         )
@@ -263,7 +285,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)
@@ -575,6 +600,7 @@ def start(self):
             self.window_inference_size = int(
                 self.window_size_choice.currentText()
             )
+            # self.window_overlap_percentage = self.window_overlap.value()
 
             self.worker = InferenceWorker(
                 device=device,
@@ -587,6 +613,7 @@ def start(self):
                 instance=self.instance_params,
                 use_window=self.use_window_inference,
                 window_infer_size=self.window_inference_size,
+                # window_overlap_percentage=self.window_overlap_percentage,
                 keep_on_cpu=self.keep_on_cpu,
                 stats_csv=self.stats_to_csv,
             )

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

@@ -113,7 +113,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]:
@@ -331,7 +330,6 @@ def fill_list_in_between(lst, n, elem):
     Returns :
         Filled list
     """
-
     new_list = []
     for i in range(len(lst)):
         temp_list = [lst[i]]
@@ -605,25 +603,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,11 +50,8 @@ install_requires =
     tifffile>=2022.2.9
     imageio-ffmpeg>=0.4.5
     torch>=1.11
-    monai>=0.9.0
-    nibabel
-    scikit-image
+    monai[nibabel,scikit-image,itk,einops]>=0.9.0
     pillow
-    itk>=5.2.0
     matplotlib
     vispy>=0.9.6