diff --git a/src/cellmap_models/pytorch/cellpose/add_model.py b/src/cellmap_models/pytorch/cellpose/add_model.py index d63d68c..d51a13d 100644 --- a/src/cellmap_models/pytorch/cellpose/add_model.py +++ b/src/cellmap_models/pytorch/cellpose/add_model.py @@ -1,3 +1,4 @@ +import os import sys from typing import Optional from cellpose.io import add_model as _add_model @@ -15,7 +16,7 @@ def add_model(model_name: Optional[str] = None): model_name = sys.argv[1] base_path = MODEL_DIR get_model(model_name, base_path) - _add_model(str(base_path / f"{model_name}.pth")) + _add_model(os.path.join(base_path, f"{model_name}.pth")) print( f"Added model {model_name}. This will now be available in the cellpose model list." ) diff --git a/src/cellmap_models/pytorch/cellpose/get_model.py b/src/cellmap_models/pytorch/cellpose/get_model.py index 0122dca..2ea7188 100644 --- a/src/cellmap_models/pytorch/cellpose/get_model.py +++ b/src/cellmap_models/pytorch/cellpose/get_model.py @@ -1,3 +1,4 @@ +import os from pathlib import Path from cellpose.utils import download_url_to_file @@ -20,10 +21,10 @@ def get_model( f"Model {model_name} is not available. Available models are {list(models_dict.keys())}." ) - if not (base_path / f"{model_name}.pth").exists(): + if not Path(base_path / f"{model_name}.pth").exists(): print(f"Downloading {model_name} from {models_dict[model_name]}") download_url_to_file( - models_dict[model_name], str(base_path / f"{model_name}.pth") + models_dict[model_name], os.path.join(base_path, f"{model_name}.pth") ) print("Downloaded model {model_name} to {base_path}.") return