WIP download weights

Author: MMathisLab

docs/res/guides/custom_model_template.rst

@@ -17,7 +17,7 @@ To add a custom model, you will need a **.py** file with the following structure
 
     def get_weights_file():
         return "weights_file.pth" # name of the weights file for the model,
-        # which should be in *napari_cellseg3d/models/saved_weights*
+        # which should be in *napari_cellseg3d/models/pretrained*
 
 
     def get_output(model, input):
@@ -35,5 +35,3 @@ To add a custom model, you will need a **.py** file with the following structure
     def ModelClass(x1,x2...):
         # your Pytorch model here...
         return results # should return as [C, N, D,H,W]
-
-

napari_cellseg3d/models/TRAILMAP_MS.py

@@ -1,10 +1,14 @@
 import torch
 from torch import nn
+from napari_cellseg3d import utils
+import os
 
+modelname = "TRAILMAP_MS"
+target_dir = os.path.join("models","pretrained")
 
 def get_weights_file():
-    # return "TMP_TEST_40e.pth"
-    return "TRAILMAP_DFl_best.pth"
+    utils.DownloadModel(model, train_dir)
+    return "TRAILMAP_MS_best_metric_epoch_26.pth" #model additionally trained on Mathis/Wyss mesoSPIM data
 
 
 def get_net():

napari_cellseg3d/models/model_SegResNet.py

@@ -1,11 +1,16 @@
 from monai.networks.nets import SegResNetVAE
+from napari_cellseg3d import utils
+import os
 
+modelname = "SegResNet"
+target_dir = os.path.join("models","pretrained")
 
 def get_net():
     return SegResNetVAE
 
 
 def get_weights_file():
+    utils.DownloadModel(modelname, target_dir)
     return "SegResNet.pth"
 
 

napari_cellseg3d/models/model_TRAILMAP.py

@@ -1,9 +1,13 @@
 from napari_cellseg3d.models.unet.model import UNet3D
+from napari_cellseg3d import utils
+import os
 
+modelname = "TRAILMAP"
+target_dir = os.path.join("models","pretrained")
 
 def get_weights_file():
-    # return "TMP_TEST_40e.pth"
-    return "trailmaptorchpretrained.pth"
+    utils.DownloadModel(model, train_dir)
+    return "TRAILMAP_PyTorch.pth" #original model form Luo lab, transfered to pytroch
 
 
 def get_net():

napari_cellseg3d/models/model_VNet.py

@@ -1,13 +1,17 @@
 from monai.inferers import sliding_window_inference
 from monai.networks.nets import VNet
+from napari_cellseg3d import utils
+import os
 
+modelname = "VNet"
+target_dir = os.path.join("models","pretrained")
 
 def get_net():
     return VNet()
 
 
 def get_weights_file():
-    # return "dice_VNet.pth"
+    utils.DownloadModel(model, train_dir)
     return "VNet_40e.pth"
 
 

napari_cellseg3d/utils.py

@@ -12,29 +12,14 @@
 from skimage.filters import gaussian
 from tifffile import imread as tfl_imread
 from tqdm import tqdm
+import importlib.util
 
 """
 utils.py
 ====================================
 Definitions of utility functions and variables
 """
 
-##################
-##################
-# dev util
-def ENABLE_TEST_MODE():
-    path = Path(os.path.expanduser("~"))
-    # print(path)
-    print("TEST MODE ENABLED, DEV ONLY")
-    if path == Path("C:/Users/Cyril"):
-        return True
-    return False
-
-
-##################
-##################
-
-
 def normalize_x(image):
     """Normalizes the values of an image array to be between [-1;1] rather than [0;255]
 
@@ -996,3 +981,64 @@ def merge_imgs(imgs, original_image_shape):
 
     print(merged_imgs.shape)
     return merged_imgs
+
+
+def read_plainconfig(configname):
+    """
+    This code is adapted from DeepLabCut with permission from MWMathis
+    """
+    if not os.path.exists(configname):
+        raise FileNotFoundError(
+            f"Config {configname} is not found. Please make sure that the file exists."
+        )
+    with open(configname) as file:
+        return YAML().load(file)
+
+def DownloadModel(modelname, target_dir):
+    """
+    Downloads a specific pretained model.
+    This code is adapted from DeepLabCut with permission from MWMathis
+    """
+    import urllib.request
+    import tarfile
+    from tqdm import tqdm
+
+    def show_progress(count, block_size, total_size):
+        pbar.update(block_size)
+
+    def tarfilenamecutting(tarf):
+        """' auxfun to extract folder path
+        ie. /xyz-trainsetxyshufflez/
+        """
+        for memberid, member in enumerate(tarf.getmembers()):
+            if memberid == 0:
+                parent = str(member.path)
+                l = len(parent) + 1
+            if member.path.startswith(parent):
+                member.path = member.path[l:]
+                yield member
+    #TODO: fix error in line 1021;
+    cellseg3d_path = os.path.split(importlib.util.find_spec("napari-cellseg3d").origin)[0]
+    neturls = read_plainconfig(os.path.join(cellseg3d_path,"models","pretrained","pretrained_model_urls.yaml",))
+
+    if modelname in neturls.keys():
+        url = neturls[modelname]
+        response = urllib.request.urlopen(url)
+        print(
+            "Downloading the model from the M.W. Mathis Lab server {}....".format(
+                url
+            )
+        )
+        total_size = int(response.getheader("Content-Length"))
+        pbar = tqdm(unit="B", total=total_size, position=0)
+        filename, _ = urllib.request.urlretrieve(url, reporthook=show_progress)
+        with tarfile.open(filename, mode="r:gz") as tar:
+            tar.extractall(target_dir, members=tarfilenamecutting(tar))
+    else:
+        models = [
+            fn
+            for fn in neturls.keys()
+            if "VNet_" not in fn and "SegResNet" not in fn and "TRAILMAP_" not in fn
+        ]
+        print("Model does not exist: ", modelname)
+        #print("Pick one of the following: ", models)