Merge pull request #19 from AdaptiveMotorControlLab/cy/log_weights_download

Moved weights download + enhancements

Author: MMathisLab

napari_cellseg3d/interface.py

@@ -2,6 +2,7 @@
 from typing import Union
 from typing import List
 
+
 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,19 +23,53 @@ 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):
+    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

@@ -15,7 +15,7 @@
 from napari_cellseg3d.models import model_SegResNet as SegResNet
 from napari_cellseg3d.models import model_TRAILMAP 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
@@ -62,8 +62,8 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
         self.models_dict = {
             "VNet": VNet,
             "SegResNet": SegResNet,
-            "TRAILMAP pre-trained": TRAILMAP,
-            "TRAILMAP_MS": TMAP,
+            "TRAILMAP": TRAILMAP,
+            "TRAILMAP_MS": TRAILMAP_MS,
         }
         """dict: dictionary of available models, with string for widget display as key
 

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: {model_name}. 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
@@ -629,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
 
@@ -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: