Repository navigation
Add CLI command to cache the cellpose model - #445
Conversation
quentinblampey
left a comment
There was a problem hiding this comment.
Thanks @alihamraoui, my comments are mostly questions, but overall the logic is great all, we can do the same for stardist!
| None, | ||
| help="Directory where the model is stored. By default, uses the cellpose default directory (`~/.cellpose/models`, or `$CELLPOSE_LOCAL_MODELS_PATH`)", | ||
| ), | ||
| gpu: bool = typer.Option(False, help="Whether to instantiate the model on GPU"), |
There was a problem hiding this comment.
Do we need the gpu argument to download the model? What happens if we always download it on CPU and then load it with a GPU for segmentation?
There was a problem hiding this comment.
models.CellposeModel takes a gpu argument to save the model on the GPU. I assume this maximizes execution speed. I think it would be useful to expose this in the CLI. do you agree?
|
|
||
| log.info(f"Cellpose model '{pretrained_model or model_type}' is cached and ready to be used.") | ||
|
|
||
| typer.echo(model_path) # stdout, so it can be captured by a workflow manager for nf-core/sopa |
There was a problem hiding this comment.
Actually, I'm wondering: since we have the model_dir option, maybe we don't need to echo the model-path? Because we already know where it's loaded?
There was a problem hiding this comment.
yes, you are right!
| def _cellpose_model_path(model) -> str: | ||
| model = getattr(model, "cp", model) | ||
|
|
||
| return model.pretrained_model |
There was a problem hiding this comment.
Does it work for both cellpose v3 and v4?
There was a problem hiding this comment.
Great question. I just saw that in Cellpose v3, model.pretrained_model returns a list of paths. I will adapt it
Thanks
Thanks @quentinblampey for the review. |
There was a problem hiding this comment.
Hi @alihamraoui, some minor comments, but everything else looks really great!
|
Thanks @quentinblampey for review, updated. |
quentinblampey
left a comment
There was a problem hiding this comment.
Awesome, thanks a lot :)
|
I see you need to fix the conflicts on the uv.lock before merging |
New
sopa cache cellposecommand to pre-download the Cellpose weights before segmentation (useful for parallel patch segmentation or offline HPC nodes).It prints the model path (by default
~/.cellpose/models) to stdout (so it can be used by the workflow manager innf-core/sopa), and--model-dirsetsCELLPOSE_LOCAL_MODELS_PATH.I factored the model initialisation out of
cellpose_patchintoload_cellpose_model, so the cache and the segmentation share the same logic.Fixes #444