fix: ⚡️ Cellpose model loading. Cosem tests.

README.md

@@ -1,4 +1,4 @@
-![GitHub Org's stars](https://img.shields.io/github/stars/Janelia-cellmap)
+[![GitHub Org's stars](https://img.shields.io/github/stars/Janelia-cellmap)](https://github.com/janelia-cellmap)
 
 
 <img src="https://raw.githubusercontent.com/janelia-cellmap/cellmap-models/main/assets/CellMapLogo2.png" alt="CellMap logo" width="85%">

src/cellmap_models/pytorch/cellpose/README.md

@@ -1,7 +1,7 @@
 <!-- FILEPATH: /Users/rhoadesj/Repos/cellmap-models/src/cellmap_models/pytorch/cellpose/README.md -->
-<h1 style="height: 56pt;">Finetuned Cellpose Models<img src="https://www.cellpose.org/static/images/cellpose_transparent.png" alt="cellpose logo" height=56pt></h1>
+<h1 style="height: 56pt;">Finetuned <a href="https://www.cellpose.org"> Cellpose Models<img src="https://www.cellpose.org/static/images/cellpose_transparent.png" alt="cellpose logo" height=56pt></h1></a>
 
-This directory contains finetuned scripts for downloading Cellpose models, particularly for use with the `cellpose` package. The models are trained on a variety of cell types from CellMap FIBSEM images, and can be used for segmentation of new data.
+This directory contains scripts for downloading finetuned Cellpose models, particularly for use with the `cellpose` package. The models are trained on a variety of cell types from CellMap FIBSEM images, and can be used for segmentation of new data. These models are returned as [CellposeModel objects](https://cellpose.readthedocs.io/en/latest/api.html#cellposemodel).
 
 ## Models
 

src/cellmap_models/pytorch/cellpose/get_model.py

@@ -28,4 +28,4 @@ def get_model(
         print(f"Downloading {model_name} from {models_dict[model_name]}")
         download_url_to_file(models_dict[model_name], full_path)
     print(f"Downloaded model {model_name} to {base_path}.")
-    return
+    return full_path

src/cellmap_models/pytorch/cellpose/load_model.py

@@ -2,12 +2,13 @@
 from pathlib import Path
 import torch
 from .get_model import get_model
+from cellpose.models import CellposeModel
 
 
 def load_model(
     model_name: str,
     base_path: str = f"{Path(__file__).parent}/models",
-    device: str = "cuda",
+    device: str | torch.device = "cuda",
 ) -> torch.nn.Module:
     """Load model
 
@@ -19,12 +20,15 @@ def load_model(
     Returns:
         model: model
     """
-
-    get_model(model_name, base_path)
+    model_path = get_model(model_name, base_path)
     if device == "cuda" and not torch.cuda.is_available():
         device = "cpu"
         print("CUDA not available. Using CPU.")
-    # TODO: load using cellpose utilities
-    model = torch.jit.load(os.path.join(base_path, f"{model_name}.pt"), device)
-    model.eval()
+    if isinstance(device, str):
+        device = torch.device(device)
+
+    model = CellposeModel(pretrained_model=model_path, device=device)
+
+    print(f"{model.diam_labels} diameter labels were used for training")
+
     return model

tests/test_cosem.py

@@ -0,0 +1,7 @@
+from cellmap_models import cosem
+
+
+def test_load_model():
+    for model_name in cosem.model_names:
+        model = cosem.load_model(model_name)
+        assert model is not None