Convert the URL file to json to save a dependency

Author: jeylau

napari_cellseg3d/models/TRAILMAP_MS.py

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

napari_cellseg3d/models/model_SegResNet.py

@@ -10,7 +10,7 @@ def get_net():
 
 
 def get_weights_file():
-    utils.DownloadModel(modelname, target_dir)
+    utils.download_model(modelname, target_dir)
     return "SegResNet.pth"
 
 

napari_cellseg3d/models/model_TRAILMAP.py

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

napari_cellseg3d/models/model_VNet.py

@@ -11,7 +11,7 @@ def get_net():
 
 
 def get_weights_file():
-    utils.DownloadModel(modelname, target_dir)
+    utils.download_model(modelname, target_dir)
     return "VNet_40e.pth"
 
 

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"
+}

napari_cellseg3d/models/pretrained/pretrained_model_urls.yaml

@@ -983,25 +983,14 @@ def merge_imgs(imgs, original_image_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):
+def download_model(modelname, target_dir):
     """
     Downloads a specific pretained model.
     This code is adapted from DeepLabCut with permission from MWMathis
     """
+    import json
     import urllib.request
     import tarfile
-    from tqdm import tqdm
 
     def show_progress(count, block_size, total_size):
         pbar.update(block_size)
@@ -1017,18 +1006,17 @@ def tarfilenamecutting(tarf):
             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",))
+    json_path = os.path.join(cellseg3d_path, "models", "pretrained", "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(
-            "Downloading the model from the M.W. Mathis Lab server {}....".format(
-                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)