Merge pull request #6 from AdaptiveMotorControlLab/mwm/download_models

MMathisLab · web-flow · commit 99412222b40a · 2022-06-28T15:22:37.000+02:00
Mwm/download models
diff --git a/docs/res/guides/custom_model_template.rst b/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]
-
-
diff --git a/napari_cellseg3d/models/TRAILMAP_MS.py b/napari_cellseg3d/models/TRAILMAP_MS.py
@@ -1,10 +1,13 @@
 import torch
 from torch import nn
+from napari_cellseg3d import utils
+import os
 
 
 def get_weights_file():
-    # return "TMP_TEST_40e.pth"
-    return "TRAILMAP_DFl_best.pth"
+    # 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")
 
 
 def get_net():
diff --git a/napari_cellseg3d/models/model_SegResNet.py b/napari_cellseg3d/models/model_SegResNet.py
@@ -1,12 +1,15 @@
 from monai.networks.nets import SegResNetVAE
+from napari_cellseg3d import utils
+import os
 
 
 def get_net():
     return SegResNetVAE
 
 
 def get_weights_file():
-    return "SegResNet.pth"
+    target_dir = utils.download_model("SegResNet")
+    return os.path.join(target_dir, "SegResNet.pth")
 
 
 def get_output(model, input):
diff --git a/napari_cellseg3d/models/model_TRAILMAP.py b/napari_cellseg3d/models/model_TRAILMAP.py
@@ -1,9 +1,12 @@
 from napari_cellseg3d.models.unet.model import UNet3D
+from napari_cellseg3d import utils
+import os
 
 
 def get_weights_file():
-    # return "TMP_TEST_40e.pth"
-    return "trailmaptorchpretrained.pth"
+    # original model from Liqun Luo lab, transfered to pytorch
+    target_dir = utils.download_model("TRAILMAP")
+    return os.path.join(target_dir, "TRAILMAP_PyTorch.pth")
 
 
 def get_net():
diff --git a/napari_cellseg3d/models/model_VNet.py b/napari_cellseg3d/models/model_VNet.py
@@ -1,14 +1,16 @@
 from monai.inferers import sliding_window_inference
 from monai.networks.nets import VNet
+from napari_cellseg3d import utils
+import os
 
 
 def get_net():
     return VNet()
 
 
 def get_weights_file():
-    # return "dice_VNet.pth"
-    return "VNet_40e.pth"
+    target_dir = utils.download_model("VNet")
+    return os.path.join(target_dir, "VNet_40e.pth")
 
 
 def get_output(model, input):
diff --git a/napari_cellseg3d/models/pretrained/pretrained_model_urls.json b/napari_cellseg3d/models/pretrained/pretrained_model_urls.json
@@ -0,0 +1,6 @@
+{
+  "TRAILMAP_MS": "http://deeplabcut.rowland.harvard.edu/cellseg3dmodels/TRAILMAP_MS.tar.gz",
+  "TRAILMAP": "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"
+}
diff --git a/napari_cellseg3d/models/pretrained/pretrained_model_urls.yaml b/napari_cellseg3d/models/pretrained/pretrained_model_urls.yaml
diff --git a/napari_cellseg3d/plugin_crop.py b/napari_cellseg3d/plugin_crop.py
@@ -72,19 +72,7 @@ def __init__(self, viewer: "napari.viewer.Viewer", parent):
 
         self.build()
 
-        ###########################################
-        if utils.ENABLE_TEST_MODE():
-            # TODO : remove/disable once done
-            if self.as_folder:
-                self.image_path = "C:/Users/Cyril/Desktop/Proj_bachelor/data/visual_png/sample"
-                if self.crop_label_choice.isChecked():
-                    self.label_path = "C:/Users/Cyril/Desktop/Proj_bachelor/data/visual_png/sample_labels"
-            else:
-                self.image_path = "C:/Users/Cyril/Desktop/Proj_bachelor/data/visual_tif/volumes/images.tif"
-                if self.crop_label_choice.isChecked():
-                    self.label_path = "C:/Users/Cyril/Desktop/Proj_bachelor/data/visual_tif/labels/testing_im.tif"
-
-        ###########################################
+        
 
     def toggle_label_path(self):
         if self.crop_label_choice.isChecked():
diff --git a/napari_cellseg3d/plugin_metrics.py b/napari_cellseg3d/plugin_metrics.py
@@ -67,23 +67,7 @@ def __init__(self, viewer: "napari.viewer.Viewer", parent):
 
         self.build()
 
-        ######################################
-        # TODO test remove
-        import glob
-        import os
-
-        if utils.ENABLE_TEST_MODE():
-            ground_directory = "C:/Users/Cyril/Desktop/Proj_bachelor/data/cropped_visual/train/lab"
-            # ground_directory = "C:/Users/Cyril/Desktop/test/labels"
-            pred_directory = "C:/Users/Cyril/Desktop/test/pred"
-            # pred_directory = "C:/Users/Cyril/Desktop/test"
-            self.images_filepaths = sorted(
-                glob.glob(os.path.join(ground_directory, "*.tif"))
-            )
-            self.labels_filepaths = sorted(
-                glob.glob(os.path.join(pred_directory, "*.tif"))
-            )
-        ###############################################################################
+        
 
     def build(self):
         """Builds the layout of the widget."""
diff --git a/napari_cellseg3d/plugin_model_training.py b/napari_cellseg3d/plugin_model_training.py
@@ -132,44 +132,7 @@ def __init__(
 
         self.save_as_zip = False
         """Whether to zip results folder once done. Creates a zipped copy of the results folder."""
-        ######################
-        ######################
-        ######################
-        # TEST TODO REMOVE
-        import glob
-
-        if utils.ENABLE_TEST_MODE():
-            directory = os.path.dirname(os.path.realpath(__file__)) + str(
-                Path("/models/dataset/volumes")
-            )
-            self.data_path = directory
-
-            lab_directory = os.path.dirname(os.path.realpath(__file__)) + str(
-                Path("/models/dataset/lab_sem")
-            )
-            self.label_path = lab_directory
-
-            self.images_filepaths = sorted(
-                glob.glob(os.path.join(directory, "*.tif"))
-            )
-
-            self.labels_filepaths = sorted(
-                glob.glob(os.path.join(lab_directory, "*.tif"))
-            )
-
-            if results_path == "":
-                self.results_path = "C:/Users/Cyril/Desktop/test/models"
-            else:
-                self.results_path = results_path
-
-            if data_path != "":
-                self.data_path = data_path
-
-            if label_path != "":
-                self.label_path = label_path
-        #######################
-        #######################
-        #######################
+        
 
         # recover default values
         self.num_samples = samples
diff --git a/napari_cellseg3d/utils.py b/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,36 @@ def merge_imgs(imgs, original_image_shape):
 
     print(merged_imgs.shape)
     return merged_imgs
+
+
+def download_model(modelname):
+    """
+    Downloads a specific pretained model.
+    This code is adapted from DeepLabCut with permission from MWMathis
+    """
+    import json
+    import urllib.request
+    import tarfile
+
+    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 modelname in neturls.keys():
+        url = neturls[modelname]
+        response = urllib.request.urlopen(url)
+        print(f"Downloading the model from the M.W. Mathis Lab server {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(pretrained_folder_path)
+        return pretrained_folder_path
+    else:
+        raise ValueError(
+            f"Unknown model. `modelname` should be one of {', '.join(neturls)}"
+        )
