Fixed model names mismatch

- TRAILMAP is now TRAILMAP_MS, as intended initially

- TRAILMAP_MS is now TRAILMAP_PyTorch.py, was incorrectly renamed

- Fixed "Pytorch" typos

- Improved logging function

- Fixed redundant imports

Co-Authored-By: Mackenzie Mathis <[email protected]>

Author: C-Achard

napari_cellseg3d/dev_scripts/weight_conversion.py

@@ -3,7 +3,7 @@
 
 import torch
 
-from napari_cellseg3d.models.model_TRAILMAP import get_net
+from napari_cellseg3d.models.model_TRAILMAP_PyTorch import get_net
 from napari_cellseg3d.models.unet.model import UNet3D
 
 # not sure this actually works when put here

napari_cellseg3d/log_utility.py

@@ -56,18 +56,20 @@ def replace_last_line(self, text):
         finally:
             self.lock.release()
 
-    def print_and_log(self, text):
+    def print_and_log(self, text, printing=True):
         """Utility used to both print to terminal and log text to a QTextEdit
          item in a thread-safe manner. Use only for important user info.
 
         Args:
             text (str): Text to be printed and logged
+            printing (bool): Whether to print the message as well or not using print(). Defaults to True.
 
         """
         self.lock.acquire()
         try:
-            print(text)
-            # causes issue if you clik on terminal (tied to CMD QuickEdit mode)
+            if printing:
+                print(text)
+            # causes issue if you clik on terminal (tied to CMD QuickEdit mode on Windows)
             self.moveCursor(QTextCursor.End)
             self.insertPlainText(f"\n{text}")
             self.verticalScrollBar().setValue(

napari_cellseg3d/model_framework.py

@@ -14,9 +14,9 @@
 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_TRAILMAP as TRAILMAP
+from napari_cellseg3d.models import model_TRAILMAP_PyTorch as TRAILMAP
 from napari_cellseg3d.models import model_VNet as VNet
-from napari_cellseg3d.models import TRAILMAP_MS as TMAP
+from napari_cellseg3d.models import model_TRAILMAP_MS as TRAILMAP_MS
 from napari_cellseg3d.plugin_base import BasePluginFolder
 
 warnings.formatwarning = utils.format_Warning
@@ -63,8 +63,8 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
         self.models_dict = {
             "VNet": VNet,
             "SegResNet": SegResNet,
-            "TRAILMAP pre-trained": TRAILMAP,
-            "TRAILMAP_MS": TMAP,
+            "TRAILMAP_PyTorch": TRAILMAP,
+            "TRAILMAP_MS": TRAILMAP_MS,
         }
         """dict: dictionary of available models, with string for widget display as key
 

napari_cellseg3d/model_workers.py

@@ -114,7 +114,7 @@ def show_progress(count, block_size, total_size):
                 tar.extractall(pretrained_folder_path)
         else:
             raise ValueError(
-                f"Unknown model. `modelname` should be one of {', '.join(neturls)}"
+                f"Unknown model: {model_name}. Should be one of {', '.join(neturls)}"
             )
 
 
@@ -695,7 +695,7 @@ def log_parameters(self):
         self.log("-" * 20)
 
     def train(self):
-        """Trains the Pytorch model for the given number of epochs, with the selected model and data,
+        """Trains the PyTorch model for the given number of epochs, with the selected model and data,
         using the chosen batch size, validation interval, loss function, and number of samples.
         Will perform validation once every :py:obj:`val_interval` and save results if the mean dice is better
 

napari_cellseg3d/models/model_SegResNet.py

@@ -1,8 +1,5 @@
-import os
-
 from monai.networks.nets import SegResNetVAE
 
-from napari_cellseg3d import utils
 
 
 def get_net():

napari_cellseg3d/models/model_TRAILMAP.py …ri_cellseg3d/models/model_TRAILMAP_MS.pynapari_cellseg3d/models/model_TRAILMAP.py renamed to napari_cellseg3d/models/model_TRAILMAP_MS.py napari_cellseg3d/models/model_TRAILMAP.py renamed to napari_cellseg3d/models/model_TRAILMAP_MS.py

@@ -1,12 +1,9 @@
-import os
-
-from napari_cellseg3d import utils
 from napari_cellseg3d.models.unet.model import UNet3D
 
 
 def get_weights_file():
-    # original model from Liqun Luo lab, transfered to pytorch
-    return "TRAILMAP_PyTorch.pth"
+    # original model from Liqun Luo lab, transferred to pytorch
+    return "TRAILMAP_MS_best_metric_epoch_26.pth"
 
 
 def get_net():

napari_cellseg3d/models/TRAILMAP_MS.py …llseg3d/models/model_TRAILMAP_PyTorch.pynapari_cellseg3d/models/TRAILMAP_MS.py renamed to napari_cellseg3d/models/model_TRAILMAP_PyTorch.py napari_cellseg3d/models/TRAILMAP_MS.py renamed to napari_cellseg3d/models/model_TRAILMAP_PyTorch.py

@@ -1,18 +1,15 @@
-import os
-
 import torch
 from torch import nn
 
-from napari_cellseg3d import utils
 
 
 def get_weights_file():
     # model additionally trained on Mathis/Wyss mesoSPIM data
-    return "TRAILMAP_MS_best_metric_epoch_26.pth"
-
+    return "TRAILMAP_PyTorch.pth"
+    # FIXME currently incorrect, find good weights from TRAILMAP_test and upload them
 
 def get_net():
-    return TRAILMAP_MS(1, 1)
+    return TRAILMAP_PyTorch(1, 1)
 
 
 def get_output(model, input):
@@ -26,7 +23,7 @@ def get_validation(model, val_inputs):
     return model(val_inputs)
 
 
-class TRAILMAP_MS(nn.Module):
+class TRAILMAP_PyTorch(nn.Module):
     def __init__(self, in_ch, out_ch):
         super().__init__()
         self.conv0 = self.encoderBlock(in_ch, 32, 3)  # input
@@ -123,3 +120,4 @@ def outBlock(self, in_ch, out_ch, kernel_size, padding="same"):
             # nn.BatchNorm3d(out_ch),
         )
         return out
+

napari_cellseg3d/models/model_VNet.py

@@ -1,9 +1,6 @@
-import os
-
 from monai.inferers import sliding_window_inference
 from monai.networks.nets import VNet
 
-from napari_cellseg3d import utils
 
 
 def get_net():

napari_cellseg3d/models/pretrained/pretrained_model_urls.json

@@ -1,6 +1,6 @@
 {
   "TRAILMAP_MS": "http://deeplabcut.rowland.harvard.edu/cellseg3dmodels/TRAILMAP_MS.tar.gz",
-  "TRAILMAP": "http://deeplabcut.rowland.harvard.edu/cellseg3dmodels/TRAILMAP.tar.gz",
+  "TRAILMAP_PyTorch": "http://deeplabcut.rowland.harvard.edu/cellseg3dmodels/TRAILMAP.tar.gz",
   "SegResNet": "http://deeplabcut.rowland.harvard.edu/cellseg3dmodels/SegResNet.tar.gz",
   "VNet": "http://deeplabcut.rowland.harvard.edu/cellseg3dmodels/VNet.tar.gz"
 }

napari_cellseg3d/plugin_model_training.py

@@ -35,7 +35,7 @@
 
 
 class Trainer(ModelFramework):
-    """A plugin to train pre-defined Pytorch models for one-channel segmentation directly in napari.
+    """A plugin to train pre-defined PyTorch models for one-channel segmentation directly in napari.
     Features parameter selection for training, dynamic loss plotting and automatic saving of the best weights during
     training through validation."""