🎨 formatting

Author: vidalmaxime

napari_cellseg3d/model_workers.py

@@ -683,7 +683,7 @@ def inference(self):
         try:
             dims = self.model_dict["segres_size"]
 
-            model = self.model_dict["class"].get_net()
+
             if self.model_dict["name"] == "SegResNet":
                 model = self.model_dict["class"].get_net()(
                     input_image_size=[
@@ -694,6 +694,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],
+                    use_checkpoint=False,
+                )
+            else:
+                model = self.model_dict["class"].get_net()
+            model = model.to(self.device)
 
             self.log_parameters()
 

napari_cellseg3d/models/model_SegResNet.py

@@ -2,7 +2,9 @@
 
 
 def get_net(input_image_size, dropout_prob=None):
-    return SegResNetVAE(input_image_size, out_channels=1, dropout_prob=dropout_prob)
+    return SegResNetVAE(
+        input_image_size, out_channels=1, dropout_prob=dropout_prob
+    )
 
 
 def get_weights_file():

napari_cellseg3d/models/model_SwinUNetR.py

@@ -7,7 +7,13 @@ def get_weights_file():
 
 
 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)
+    return SwinUNETR(
+        img_size,
+        in_channels=1,
+        out_channels=1,
+        feature_size=48,
+        use_checkpoint=use_checkpoint,
+    )
 
 
 def get_output(model, input):