fix: 📦️ Remove unused files and update model loading

.gitignore

@@ -157,4 +157,7 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+#.idea/
+
+# Pytorch
+*.pth

README.md

@@ -20,13 +20,17 @@ pip install .
 ## Usage
 
 ```python
-import cellmap-models
+import cellmap_models
 ```
 
 Different models are available in the `cellmap-models` module. For example, to use the models produced by the `COSEM` pilot project team, and published as part of [Whole-cell organelle segmentation in volume electron microscopy](https://doi.org/10.1038/s41586-021-03977-3):
 
 ```python
-import cellmap-models.cosem as cosem_models
+import cellmap_models.cosem as cosem_models
+print(cosem_models.models_list)
+```
+This will list the available models. To load a specific model, use the `load_model` function:
+```python
 model = cosem_models.load_model('setup04/1820500')
 ```
 

src/cellmap_models/__init__.py

@@ -3,3 +3,4 @@
 """
 
 from .utils import download_url_to_file
+from .pytorch import cosem

src/cellmap_models/pytorch/__init__.py

@@ -0,0 +1 @@
+from . import cosem

src/cellmap_models/pytorch/cosem/README.md

@@ -4,26 +4,24 @@ This repository contains the COSEM trained networks, converted to PyTorch. The o
 
 The networks have been converted to PyTorch from their original Tensorflow versions using the scripts available [here](https://github.com/pattonw/cnnectome.conversion). All models are trained on 3D data and expect input of shape `(batch_size, 1, z, y, x)`.
 
-This repo is pip installable, simply follow the steps below in an appropriate python environment (python >= 3.7), replacing `/path/to/cosem_models` with the path to this repository on your local machine:
+You can load a model using the following code:
 
-```bash
-cd /path/to/cosem_models
-pip install .
+```python
+import cellmap_models.cosem as cosem_models
+model = cosem_models.load_model('setup04/1820500')
+# The model is now ready to use
 ```
 
-Then you can load a model using the following code:
+Available models can be seen with the following code:
 
 ```python
-import cosem_models
-model = cosem_models.load_model('setup04/1820500')
-
-# The model is now ready to use
+cosem_models.models_list
 ```
 
 Each model has a separate unet backbone and single layer prediction head. The `unet` and `head` objects are both PyTorch modules and can be used as such. You can access the separate components of the model using the following code:
 
 ```python
-import cosem_models
+import cellmap_models.cosem as cosem_models
 model = cosem_models.load_model('setup04/1820500')
 unet = model.unet
 head = model.prediction_head

src/cellmap_models/pytorch/cosem/__init__.py

@@ -0,0 +1,14 @@
+from .load_model import load_model
+models_list = [
+    "setup04/1820500.pth",
+    "setup04/625000.pth",
+    "setup04/975000.pth",
+    "setup26.1/2580000.pth",
+    "setup26.1/650000.pth",
+    "setup28/1440000.pth",
+    "setup28/775000.pth",
+    "setup36/1100000.pth",
+    "setup36/500000.pth",
+    "setup45/1634500.pth",
+    "setup45/625000.pth",
+]

src/cellmap_models/pytorch/cosem/load_model.py

@@ -30,6 +30,8 @@ def get_param_dict(model_params):
 
 
 def load_model(checkpoint_path):
+    if not Path(checkpoint_path).exists():
+        checkpoint_path = Path(Path(__file__).parent / checkpoint_path)
     model_params = SourceFileLoader(
         "model", str(Path(checkpoint_path).parent / "model.py")
     ).load_module()

src/cellmap_models/pytorch/cosem/setup04/model.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 import numpy as np
-from cellmap-models import download_url_to_file
+from cellmap_models import download_url_to_file
 
 # voxel size parameters
 voxel_size_output = np.array((4,) * 3)
@@ -34,4 +34,6 @@
     "975000": "https://janelia-cosem-networks.s3.amazonaws.com/v0003.2-pytorch/cosem_models/cosem_models/setup04/975000",
 }
 for name, url in urls.items():
-    download_url_to_file(url, str(Path(__file__).parent / f"{name}"))
+    if not (Path(__file__).parent / f"{name}").exists():
+        print(f"Downloading {name} from {url}")
+        download_url_to_file(url, str(Path(__file__).parent / f"{name}.pth"))

src/cellmap_models/pytorch/cosem/setup26.1/model.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 import numpy as np
-from cellmap-models import download_url_to_file
+from cellmap_models import download_url_to_file
 
 # voxel size parameters
 voxel_size_output = np.array((4,) * 3)
@@ -33,4 +33,6 @@
     "650000": "https://janelia-cosem-networks.s3.amazonaws.com/v0003.2-pytorch/cosem_models/cosem_models/setup26.1/650000",
 }
 for name, url in urls.items():
-    download_url_to_file(url, str(Path(__file__).parent / f"{name}"))
+    if not (Path(__file__).parent / f"{name}").exists():
+        print(f"Downloading {name} from {url}")
+        download_url_to_file(url, str(Path(__file__).parent / f"{name}.pth"))

src/cellmap_models/pytorch/cosem/setup28/model.py

@@ -1,6 +1,6 @@
 import numpy as np
 from pathlib import Path
-from cellmap-models import download_url_to_file
+from cellmap_models import download_url_to_file
 
 # voxel size parameters
 voxel_size_output = np.array((4,) * 3)
@@ -33,4 +33,6 @@
     "775000": "https://janelia-cosem-networks.s3.amazonaws.com/v0003.2-pytorch/cosem_models/cosem_models/setup28/775000",
 }
 for name, url in urls.items():
-    download_url_to_file(url, str(Path(__file__).parent / f"{name}"))
+    if not (Path(__file__).parent / f"{name}").exists():
+        print(f"Downloading {name} from {url}")
+        download_url_to_file(url, str(Path(__file__).parent / f"{name}.pth"))

src/cellmap_models/pytorch/cosem/setup36/model.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 import numpy as np
-from cellmap-models import download_url_to_file
+from cellmap_models import download_url_to_file
 
 # voxel size parameters
 voxel_size_output = np.array((4,) * 3)
@@ -33,4 +33,7 @@
     "500000": "https://janelia-cosem-networks.s3.amazonaws.com/v0003.2-pytorch/cosem_models/cosem_models/setup36/500000",
 }
 for name, url in urls.items():
-    download_url_to_file(url, str(Path(__file__).parent / f"{name}"))
+    if not (Path(__file__).parent / f"{name}").exists():
+        print(f"Downloading {name} from {url}")
+        download_url_to_file(url, str(Path(__file__).parent / f"{name}.pth"))
+

src/cellmap_models/pytorch/cosem/setup45/model.py

@@ -1,6 +1,6 @@
 import numpy as np
 from pathlib import Path
-from cellmap-models import download_url_to_file
+from cellmap_models import download_url_to_file
 
 # voxel size parameters
 voxel_size_output = np.array((4,) * 3)
@@ -33,4 +33,7 @@
     "625000": "https://janelia-cosem-networks.s3.amazonaws.com/v0003.2-pytorch/cosem_models/cosem_models/setup45/625000",
 }
 for name, url in urls.items():
-    download_url_to_file(url, str(Path(__file__).parent / f"{name}"))
+    if not (Path(__file__).parent / f"{name}").exists():
+        print(f"Downloading {name} from {url}")
+        download_url_to_file(url, str(Path(__file__).parent / f"{name}.pth"))
+