Moved weights download + enhancements

- Weight download now always happens in worker thread to avoid freezing UI

- Download progress bar now shows on the log in the plugin

- Models .py files changed to be as simple as possible (so users can more easily add custom models)

Author: C-Achard

.gitignore

@@ -10,6 +10,7 @@ __pycache__/
 *.tif
 napari_cellseg3d/_tests/res/*.csv
 *.pth
+*.db
 
 # Distribution / packaging
 .Python

napari_cellseg3d/interface.py

@@ -1,6 +1,7 @@
 from typing import Optional
 from typing import Union
 
+
 from qtpy.QtCore import Qt
 from qtpy.QtCore import QUrl
 from qtpy.QtGui import QDesktopServices

napari_cellseg3d/log_utility.py

@@ -1,5 +1,6 @@
 import threading
 
+from qtpy import QtCore
 from qtpy.QtGui import QTextCursor
 from qtpy.QtWidgets import QTextEdit
 
@@ -22,6 +23,38 @@ def __init__(self, parent):
 
     # def receive_log(self, text):
     #     self.print_and_log(text)
+    def write(self, message):
+        self.lock.acquire()
+        try:
+            if not hasattr(self, "flag"):
+                self.flag = False
+            message = message.replace('\r', '').rstrip()
+            if message:
+                method = "replace_last_line" if self.flag else "append"
+                QtCore.QMetaObject.invokeMethod(self,
+                                                method,
+                                                QtCore.Qt.QueuedConnection,
+                                                QtCore.Q_ARG(str, message))
+                self.flag = True
+            else:
+                self.flag = False
+
+        finally:
+            self.lock.release()
+
+    @QtCore.Slot(str)
+    def replace_last_line(self, text):
+        self.lock.acquire()
+        try:
+            cursor = self.textCursor()
+            cursor.movePosition(QTextCursor.End)
+            cursor.select(QTextCursor.BlockUnderCursor)
+            cursor.removeSelectedText()
+            cursor.insertBlock()
+            self.setTextCursor(cursor)
+            self.insertPlainText(text)
+        finally:
+            self.lock.release()
 
     def print_and_log(self, text):
         """Utility used to both print to terminal and log text to a QTextEdit

napari_cellseg3d/model_workers.py

@@ -1,9 +1,13 @@
 import os
 import platform
 from pathlib import Path
+import importlib.util
+from typing import Optional
 
 import numpy as np
+from tifffile import imwrite
 import torch
+from tqdm import tqdm
 
 # MONAI
 from monai.data import CacheDataset
@@ -37,9 +41,10 @@
 
 # Qt
 from qtpy.QtCore import Signal
-from tifffile import imwrite
+
 
 from napari_cellseg3d import utils
+from napari_cellseg3d import log_utility
 
 # local
 from napari_cellseg3d.model_instance_seg import binary_connected
@@ -57,9 +62,61 @@
 # https://napari-staging-site.github.io/guides/stable/threading.html
 
 WEIGHTS_DIR = os.path.dirname(os.path.realpath(__file__)) + str(
-    Path("/models/saved_weights")
+    Path("/models/pretrained")
 )
 
+class WeightsDownloader:
+
+    def __init__(self, log_widget: Optional[log_utility.Log]= None):
+        self.log_widget = log_widget
+
+    def download_weights(self,model_name: str):
+        """
+            Downloads a specific pretrained model.
+            This code is adapted from DeepLabCut with permission from MWMathis.
+
+            Args:
+                model_name (str): name of the model to download
+            """
+        import json
+        import tarfile
+        import urllib.request
+
+        def show_progress(count, block_size, total_size):
+            pbar.update(block_size)
+
+        cellseg3d_path = os.path.split(
+            importlib.util.find_spec("napari_cellseg3d").origin
+        )[0]
+        pretrained_folder_path = os.path.join(
+            cellseg3d_path, "models", "pretrained"
+        )
+        json_path = os.path.join(
+            pretrained_folder_path, "pretrained_model_urls.json"
+        )
+        with open(json_path) as f:
+            neturls = json.load(f)
+        if model_name in neturls.keys():
+            url = neturls[model_name]
+            response = urllib.request.urlopen(url)
+
+            start_message = f"Downloading the model from the M.W. Mathis Lab server {url}...."
+            total_size = int(response.getheader("Content-Length"))
+            if self.log_widget is None:
+                print(start_message)
+                pbar = tqdm(unit="B", total=total_size, position=0)
+            else:
+                self.log_widget.print_and_log(start_message)
+                pbar = tqdm(unit="B", total=total_size, position=0, file=self.log_widget)
+
+            filename, _ = urllib.request.urlretrieve(url, reporthook=show_progress)
+            with tarfile.open(filename, mode="r:gz") as tar:
+                tar.extractall(pretrained_folder_path)
+        else:
+            raise ValueError(
+                f"Unknown model. `modelname` should be one of {', '.join(neturls)}"
+            )
+
 
 class LogSignal(WorkerBaseSignals):
     """Signal to send messages to be logged from another thread.
@@ -142,9 +199,12 @@ def __init__(
         self.window_infer_size = window_infer_size
         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"""
 
+        self.downloader = WeightsDownloader()
+        """Download utility"""
+
+
     @staticmethod
     def create_inference_dict(images_filepaths):
         """Create a dict for MONAI with "image" keys with all image paths in :py:attr:`~self.images_filepaths`
@@ -154,6 +214,9 @@ def create_inference_dict(images_filepaths):
         data_dicts = [{"image": image_name} for image_name in images_filepaths]
         return data_dicts
 
+    def set_download_log(self, widget):
+        self.downloader.log_widget = widget
+
     def log(self, text):
         """Sends a signal that ``text`` should be logged
 
@@ -233,8 +296,8 @@ def inference(self):
         """
         sys = platform.system()
         print(f"OS is {sys}")
-        if sys == "Darwin":  # required for macOS ?
-            torch.set_num_threads(1)
+        if sys == "Darwin":
+            torch.set_num_threads(1) # required for threading on macOS ?
             self.log("Number of threads has been set to 1 for macOS")
 
         images_dict = self.create_inference_dict(self.images_filepaths)
@@ -260,7 +323,7 @@ def inference(self):
         model = self.model_dict["class"].get_net()
         if self.model_dict["name"] == "SegResNet":
             model = self.model_dict["class"].get_net()(
-                input_image_size=[dims, dims, dims],  # TODO FIX !
+                input_image_size=[dims, dims, dims],  # TODO FIX ! find a better way & remove model-specific code
                 out_channels=1,
                 # dropout_prob=0.3,
             )
@@ -304,12 +367,13 @@ def inference(self):
         # print(weights)
         self.log(
             "\nLoading weights..."
-        )  # TODO add try/except for invalid weights
+        )  # TODO add try/except for invalid weights for proper reset
 
         if self.weights_dict["custom"]:
             weights = self.weights_dict["path"]
         else:
-            weights = os.path.join(WEIGHTS_DIR, self.weights_dict["path"])
+            self.downloader.download_weights(self.model_dict["name"])
+            weights = os.path.join(WEIGHTS_DIR, self.model_dict["class"].get_weights_file())
 
         model.load_state_dict(
             torch.load(
@@ -544,8 +608,6 @@ def __init__(
 
 
         """
-
-        print("init")
         super().__init__(self.train)
         self._signals = LogSignal()
         self.log_signal = self._signals.log_signal
@@ -571,7 +633,11 @@ def __init__(
 
         self.train_files = []
         self.val_files = []
-        print("end init")
+
+        self.downloader = WeightsDownloader()
+
+    def set_download_log(self, widget):
+        self.downloader.log_widget = widget
 
     def log(self, text):
         """Sends a signal that ``text`` should be logged
@@ -838,6 +904,7 @@ def train(self):
         if self.weights_path is not None:
             if self.weights_path == "use_pretrained":
                 weights_file = model_class.get_weights_file()
+                self.downloader.download_weights(model_name)
                 weights = os.path.join(WEIGHTS_DIR, weights_file)
                 self.weights_path = weights
             else:

napari_cellseg3d/models/TRAILMAP_MS.py

@@ -8,8 +8,7 @@
 
 def get_weights_file():
     # model additionally trained on Mathis/Wyss mesoSPIM data
-    target_dir = utils.download_model("TRAILMAP_MS")
-    return os.path.join(target_dir, "TRAILMAP_MS_best_metric_epoch_26.pth")
+    return "TRAILMAP_MS_best_metric_epoch_26.pth"
 
 
 def get_net():

napari_cellseg3d/models/model_SegResNet.py

@@ -10,8 +10,7 @@ def get_net():
 
 
 def get_weights_file():
-    target_dir = utils.download_model("SegResNet")
-    return os.path.join(target_dir, "SegResNet.pth")
+    return "SegResNet.pth"
 
 
 def get_output(model, input):

napari_cellseg3d/models/model_TRAILMAP.py

@@ -6,8 +6,7 @@
 
 def get_weights_file():
     # original model from Liqun Luo lab, transfered to pytorch
-    target_dir = utils.download_model("TRAILMAP")
-    return os.path.join(target_dir, "TRAILMAP_PyTorch.pth")
+    return "TRAILMAP_PyTorch.pth"
 
 
 def get_net():

napari_cellseg3d/models/model_VNet.py

@@ -11,8 +11,7 @@ def get_net():
 
 
 def get_weights_file():
-    target_dir = utils.download_model("VNet")
-    return os.path.join(target_dir, "VNet_40e.pth")
+    return "VNet_40e.pth"
 
 
 def get_output(model, input):

napari_cellseg3d/plugin_model_inference.py

@@ -534,7 +534,6 @@ def start(self):
             else:
                 weights_dict = {
                     "custom": False,
-                    "path": self.get_model(model_key).get_weights_file(),
                 }
 
             if self.anisotropy_wdgt.is_enabled():
@@ -591,6 +590,7 @@ def start(self):
                 keep_on_cpu=self.keep_on_cpu,
                 stats_csv=self.stats_to_csv,
             )
+            self.worker.set_download_log(self.log)
 
             yield_connect_show_res = lambda data: self.on_yield(
                 data,

napari_cellseg3d/plugin_model_training.py

@@ -851,6 +851,7 @@ def start(self):
                 do_augmentation=self.augment_choice.isChecked(),
                 deterministic=seed_dict,
             )
+            self.worker.set_download_log(self.log)
 
             [btn.setVisible(False) for btn in self.close_buttons]
 
@@ -988,6 +989,11 @@ def on_yield(data, widget):
     def make_csv(self):
 
         size_column = range(1, self.max_epochs + 1)
+
+        if len(self.loss_values) == 0 or self.loss_values is None:
+            warnings.warn("No loss values to add to csv !")
+            return
+        
         self.df = pd.DataFrame(
             {
                 "epoch": size_column,