Fixed model name error

- Replaced all TRAILMAP_PyTorch with TRAILMAP

Author: C-Achard

napari_cellseg3d/dev_scripts/weight_conversion.py

@@ -3,7 +3,7 @@
 
 import torch
 
-from napari_cellseg3d.models.model_TRAILMAP_PyTorch import get_net
+from napari_cellseg3d.models.model_TRAILMAP import get_net
 from napari_cellseg3d.models.unet.model import UNet3D
 
 # not sure this actually works when put here

napari_cellseg3d/model_framework.py

@@ -14,7 +14,7 @@
 from napari_cellseg3d import utils
 from napari_cellseg3d.log_utility import Log
 from napari_cellseg3d.models import model_SegResNet as SegResNet
-from napari_cellseg3d.models import model_TRAILMAP_PyTorch as TRAILMAP
+from napari_cellseg3d.models import model_TRAILMAP as TRAILMAP
 from napari_cellseg3d.models import model_VNet as VNet
 from napari_cellseg3d.models import model_TRAILMAP_MS as TRAILMAP_MS
 from napari_cellseg3d.plugin_base import BasePluginFolder
@@ -63,7 +63,7 @@ def __init__(self, viewer: "napari.viewer.Viewer"):
         self.models_dict = {
             "VNet": VNet,
             "SegResNet": SegResNet,
-            "TRAILMAP_PyTorch": TRAILMAP,
+            "TRAILMAP": TRAILMAP,
             "TRAILMAP_MS": TRAILMAP_MS,
         }
         """dict: dictionary of available models, with string for widget display as key

napari_cellseg3d/models/model_TRAILMAP_PyTorch.py

@@ -5,11 +5,11 @@
 
 def get_weights_file():
     # model additionally trained on Mathis/Wyss mesoSPIM data
-    return "TRAILMAP_PyTorch.pth"
+    return "TRAILMAP.pth"
     # FIXME currently incorrect, find good weights from TRAILMAP_test and upload them
 
 def get_net():
-    return TRAILMAP_PyTorch(1, 1)
+    return TRAILMAP(1, 1)
 
 
 def get_output(model, input):
@@ -23,7 +23,7 @@ def get_validation(model, val_inputs):
     return model(val_inputs)
 
 
-class TRAILMAP_PyTorch(nn.Module):
+class TRAILMAP(nn.Module):
     def __init__(self, in_ch, out_ch):
         super().__init__()
         self.conv0 = self.encoderBlock(in_ch, 32, 3)  # input

napari_cellseg3d/models/pretrained/pretrained_model_urls.json

@@ -1,6 +1,6 @@
 {
   "TRAILMAP_MS": "http://deeplabcut.rowland.harvard.edu/cellseg3dmodels/TRAILMAP_MS.tar.gz",
-  "TRAILMAP_PyTorch": "http://deeplabcut.rowland.harvard.edu/cellseg3dmodels/TRAILMAP.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"
 }