Skip to content

Add CLI command to cache the cellpose model - #445

Merged
alihamraoui merged 9 commits into
mainfrom
cache_pre_trained_models
Sep 2, 2026
Merged

alihamraoui merged 9 commits into
mainfrom
cache_pre_trained_models

Conversation

@alihamraoui

@alihamraoui alihamraoui commented Aug 27, 2026 •

Copy link
Copy Markdown
Collaborator

New sopa cache cellpose command 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 in nf-core/sopa ), and --model-dir sets CELLPOSE_LOCAL_MODELS_PATH.

I factored the model initialisation out of cellpose_patch into load_cellpose_model, so the cache and the segmentation share the same logic.

Fixes #444

@quentinblampey quentinblampey left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @alihamraoui, my comments are mostly questions, but overall the logic is great all, we can do the same for stardist!

Comment thread sopa/cli/cache.py Outdated
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"),

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Comment thread sopa/cli/cache.py Outdated

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, you are right!

Comment thread sopa/cli/cache.py Outdated
def _cellpose_model_path(model) -> str:
model = getattr(model, "cp", model)

return model.pretrained_model

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does it work for both cellpose v3 and v4?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great question. I just saw that in Cellpose v3, model.pretrained_model returns a list of paths. I will adapt it
Thanks

@alihamraoui

Copy link
Copy Markdown
Collaborator Author

Thanks @alihamraoui, my comments are mostly questions, but overall the logic is great all, we can do the same for stardist!

Thanks @quentinblampey for the review.

@quentinblampey quentinblampey left a comment •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @alihamraoui, some minor comments, but everything else looks really great!

Comment thread sopa/segmentation/methods/_cellpose.py
Comment thread sopa/segmentation/methods/_cellpose.py
Comment thread sopa/cli/download.py Outdated
@alihamraoui

Copy link
Copy Markdown
Collaborator Author

Thanks @quentinblampey for review, updated.

@quentinblampey quentinblampey left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Awesome, thanks a lot :)

@quentinblampey

Copy link
Copy Markdown
Collaborator

I see you need to fix the conflicts on the uv.lock before merging
Advice: you can merge main to cache_pre_trained_models locally, keep all changes from main and push again on cache_pre_trained_models

@alihamraoui
alihamraoui merged commit 918aedb into main Sep 2, 2026
20 checks passed
@alihamraoui
alihamraoui deleted the cache_pre_trained_models branch September 2, 2026 20:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature] CLI command to cache the cellpose and/or stardist model

2 participants