From b7baf7d39d0dac6e078279452871feb243da6a41 Mon Sep 17 00:00:00 2001 From: Arjun Raj Date: Mon, 29 Jun 2026 21:11:03 -0400 Subject: [PATCH 1/2] cellposesam: default to cpsam_v2 checkpoint, keep legacy cpsam Cellpose 4.2.x ships a new Cellpose-SAM checkpoint, cpsam_v2 (SAM-ViTL, June 2026), which reduces spurious masks in low-contrast regions vs the original April 2025 cpsam. Make it the default while keeping the old model selectable for reproducibility. - environment.yml: cellpose 4.0.1 -> 4.2.1.1 (cpsam_v2 lands in 4.2.x) - models_config.py: new import-free single source of truth mapping dropdown labels -> checkpoints ('cellpose-sam' -> cpsam_v2, 'cellpose-sam (legacy cpsam)' -> cpsam). Unit-tested in the local venv. - entrypoint.py: pass pretrained_model explicitly for base models so behavior is pinned to the selection, not cellpose's internal default (which can shift between versions). Default stays the string 'cellpose-sam', so saved tool configs keep working (now -> cpsam_v2). - download_models.py: pre-download both checkpoints at build time. - Dockerfile: copy models_config.py before the download step. - CELLPOSESAM.md: document the new default, legacy option, and version. Verified locally: 4/4 mapping unit tests pass; all worker .py files compile. The CUDA build + GPU inference smoke test must be run on an amd64/GPU host (cannot run in this environment). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_015xrXEVpvb4c4ScEVjN1VsM --- .../annotations/cellposesam/CELLPOSESAM.md | 10 +++-- workers/annotations/cellposesam/Dockerfile | 3 ++ .../cellposesam/download_models.py | 14 ++++++- workers/annotations/cellposesam/entrypoint.py | 15 +++++-- .../annotations/cellposesam/environment.yml | 2 +- .../annotations/cellposesam/models_config.py | 26 ++++++++++++ .../annotations/cellposesam/tests/__init__.py | 0 .../cellposesam/tests/test_models_config.py | 41 +++++++++++++++++++ 8 files changed, 102 insertions(+), 9 deletions(-) create mode 100644 workers/annotations/cellposesam/models_config.py create mode 100644 workers/annotations/cellposesam/tests/__init__.py create mode 100644 workers/annotations/cellposesam/tests/test_models_config.py diff --git a/workers/annotations/cellposesam/CELLPOSESAM.md b/workers/annotations/cellposesam/CELLPOSESAM.md index 5297bbb..f94ca27 100644 --- a/workers/annotations/cellposesam/CELLPOSESAM.md +++ b/workers/annotations/cellposesam/CELLPOSESAM.md @@ -5,7 +5,7 @@ This worker runs Cellpose-SAM, a variant of Cellpose that combines Cellpose with ## How It Works 1. **Channel Assembly**: Collects up to three input channels from user-selected channel checkboxes and stacks them -2. **Model Selection**: Loads the built-in `cellpose-sam` model or a user-trained model from Girder +2. **Model Selection**: Loads a built-in Cellpose-SAM checkpoint (`cpsam_v2` by default, or the original `cpsam`) or a user-trained model from Girder 3. **Tiling**: Splits the image into overlapping tiles using DeepTile 4. **Segmentation**: Runs Cellpose-SAM inference on each tile with GPU acceleration 5. **Stitching**: Merges polygons spanning tile boundaries using DeepTile's `stitch_polygons()` @@ -19,7 +19,7 @@ This worker runs Cellpose-SAM, a variant of Cellpose that combines Cellpose with | **Batch XY** | text | -- | XY positions to iterate over (e.g., "1-3, 5-8") | | **Batch Z** | text | -- | Z slices to iterate over | | **Batch Time** | text | -- | Time points to iterate over | -| **Model** | select | cellpose-sam | Model to use. Includes built-in `cellpose-sam` and user-trained models from Girder | +| **Model** | select | cellpose-sam | Model to use. `cellpose-sam` runs the `cpsam_v2` checkpoint (current default); `cellpose-sam (legacy cpsam)` runs the original April 2025 `cpsam` checkpoint. User-trained models from Girder are also listed | | **Channel for Slot 1** | channelCheckboxes | -- | **Required.** Source channel(s) for the model's first input slot. If multiple selected, only the first is used | | **Channel for Slot 2** | channelCheckboxes | -- | Optional second input slot channel | | **Channel for Slot 3** | channelCheckboxes | -- | Optional third input slot channel | @@ -42,9 +42,13 @@ Unlike the standard Cellpose worker which uses Primary/Secondary channel selecto ### Model Behavior -- **Base model** (`cellpose-sam`): Runs with `gpu=True` and no diameter or channel parameters in `eval_parameters`, relying on the model's built-in defaults +- **Base models**: The dropdown labels map to cellpose built-in checkpoints in `models_config.py` — `cellpose-sam` → `cpsam_v2`, `cellpose-sam (legacy cpsam)` → `cpsam`. The selected checkpoint name is passed explicitly as `pretrained_model` (rather than relying on cellpose's internal default, which can shift between versions). Runs with `gpu=True` and no diameter/channel parameters in `eval_parameters`. - **Custom models**: Loaded from Girder by path and use the user-specified diameter in `eval_parameters` +### Built-in Checkpoints + +`cpsam_v2` (SAM-ViTL backbone, released June 2026) is the default; it reduces spurious masks in low-contrast regions compared to the original `cpsam` (April 2025). Both checkpoints (~1.23 GB each) are pre-downloaded at build time by `download_models.py` so neither downloads on first run. Requires `cellpose==4.2.1.1` (cpsam_v2 was added in the 4.2.x line). To add or change the offered checkpoints, edit `models_config.py` — it is the single source of truth for both the interface and the build-time download. + ### GPU Handling The worker always requests GPU mode (`gpu=True`). Cellpose handles the fallback to CPU internally if no GPU is available. diff --git a/workers/annotations/cellposesam/Dockerfile b/workers/annotations/cellposesam/Dockerfile index ad2ccf0..e0b5b02 100644 --- a/workers/annotations/cellposesam/Dockerfile +++ b/workers/annotations/cellposesam/Dockerfile @@ -45,6 +45,9 @@ RUN git clone https://github.com/Kitware/UPennContrast/ RUN pip install -r /UPennContrast/devops/girder/annotation_client/requirements.txt RUN pip install -e /UPennContrast/devops/girder/annotation_client/ +# models_config.py is the single source of truth for the built-in checkpoints, +# so it must be present before download_models.py runs. +COPY ./workers/annotations/cellposesam/models_config.py / COPY ./workers/annotations/cellposesam/download_models.py / RUN python /download_models.py diff --git a/workers/annotations/cellposesam/download_models.py b/workers/annotations/cellposesam/download_models.py index 255ee47..07c0082 100644 --- a/workers/annotations/cellposesam/download_models.py +++ b/workers/annotations/cellposesam/download_models.py @@ -1,3 +1,15 @@ +"""Pre-download the built-in Cellpose-SAM checkpoints at build time. + +Instantiating ``CellposeModel(pretrained_model=)`` fetches the weights to +``models.MODELS_DIR`` on first use, so doing it here bakes them into the image +and avoids a multi-GB download on the first run. ``gpu=True`` matches runtime; +on a GPU-less build host cellpose falls back to CPU but still downloads. +""" + from cellpose import models -model = models.CellposeModel(gpu=True) +from models_config import BASE_MODEL_CHECKPOINTS + +for checkpoint in set(BASE_MODEL_CHECKPOINTS.values()): + print(f"Downloading Cellpose-SAM checkpoint: {checkpoint}") + models.CellposeModel(gpu=True, pretrained_model=checkpoint) diff --git a/workers/annotations/cellposesam/entrypoint.py b/workers/annotations/cellposesam/entrypoint.py index 4933cea..26ef9ed 100644 --- a/workers/annotations/cellposesam/entrypoint.py +++ b/workers/annotations/cellposesam/entrypoint.py @@ -14,7 +14,7 @@ from worker_client import WorkerClient, geometry_to_polygon_coords -BASE_MODELS = ['cellpose-sam'] +from models_config import BASE_MODELS, BASE_MODEL_CHECKPOINTS, DEFAULT_MODEL def interface(image, apiUrl, token): @@ -68,8 +68,10 @@ def interface(image, apiUrl, token): 'Model': { 'type': 'select', 'items': models, - 'default': 'cellpose-sam', - 'tooltip': 'cellpose-sam is the base model', + 'default': DEFAULT_MODEL, + 'tooltip': 'cellpose-sam runs the cpsam_v2 checkpoint (the current default).\n' + 'Choose "cellpose-sam (legacy cpsam)" to reproduce results from the ' + 'original April 2025 model. Custom trained models are also listed here.', 'noCache': True, 'displayOrder': 4, }, @@ -249,8 +251,13 @@ def compute(datasetId, apiUrl, token, params): print(f"Models directory contents: {list(MODELS_DIR.glob('*'))}") if model in BASE_MODELS: + # Pass the checkpoint name explicitly so behavior is pinned to the + # selected model rather than relying on cellpose's internal default, + # which can change between cellpose versions. + checkpoint = BASE_MODEL_CHECKPOINTS[model] cellpose = cellpose_segmentation( - model_parameters={'gpu': True}, eval_parameters={}, output_format='polygons') + model_parameters={'gpu': True, 'pretrained_model': checkpoint}, + eval_parameters={}, output_format='polygons') else: # Get the full path to the model model_path = str(MODELS_DIR / model) diff --git a/workers/annotations/cellposesam/environment.yml b/workers/annotations/cellposesam/environment.yml index 408ba05..e15dd68 100644 --- a/workers/annotations/cellposesam/environment.yml +++ b/workers/annotations/cellposesam/environment.yml @@ -6,5 +6,5 @@ dependencies: - rasterio - shapely - pip: - - cellpose==4.0.1 + - cellpose==4.2.1.1 - deeptile \ No newline at end of file diff --git a/workers/annotations/cellposesam/models_config.py b/workers/annotations/cellposesam/models_config.py new file mode 100644 index 0000000..df950af --- /dev/null +++ b/workers/annotations/cellposesam/models_config.py @@ -0,0 +1,26 @@ +"""Built-in Cellpose-SAM model options for the cellposesam worker. + +Maps the human-friendly names shown in the Model dropdown to the cellpose +checkpoint identifiers passed to ``CellposeModel(pretrained_model=...)``. + +Kept deliberately import-free (no cellpose/deeptile/annotation_client) so the +mapping can be unit-tested in the lightweight local venv without the full +worker stack. +""" + +# Human-friendly dropdown label -> cellpose built-in checkpoint name. +# Insertion order lists the default (cpsam_v2) first. +BASE_MODEL_CHECKPOINTS = { + # New default: SAM-ViTL backbone, released June 2026; fixes spurious masks + # in low-contrast regions relative to the original cpsam. + 'cellpose-sam': 'cpsam_v2', + # Original Cellpose-SAM model, released April 2025. Kept selectable so prior + # results remain reproducible. + 'cellpose-sam (legacy cpsam)': 'cpsam', +} + +# The dropdown option selected by default. +DEFAULT_MODEL = 'cellpose-sam' + +# The list of built-in model labels offered in the interface. +BASE_MODELS = list(BASE_MODEL_CHECKPOINTS) diff --git a/workers/annotations/cellposesam/tests/__init__.py b/workers/annotations/cellposesam/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/workers/annotations/cellposesam/tests/test_models_config.py b/workers/annotations/cellposesam/tests/test_models_config.py new file mode 100644 index 0000000..15b17ef --- /dev/null +++ b/workers/annotations/cellposesam/tests/test_models_config.py @@ -0,0 +1,41 @@ +"""Unit tests for the built-in Cellpose-SAM model mapping. + +These exercise ``models_config`` in isolation — it must stay free of heavy +imports (cellpose/deeptile/annotation_client) so it runs in the lightweight +local venv without the full worker stack. Run with: + + .cache/testvenv/bin/pytest workers/annotations/cellposesam/tests -q +""" + +import sys +from pathlib import Path + +# Put the worker directory (parent of tests/) on the path so we can import the +# standalone mapping module without installing the whole worker. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +import models_config # noqa: E402 + + +def test_default_model_resolves_to_cpsam_v2(): + """The default dropdown selection must run the new cpsam_v2 checkpoint.""" + checkpoint = models_config.BASE_MODEL_CHECKPOINTS[models_config.DEFAULT_MODEL] + assert checkpoint == 'cpsam_v2' + + +def test_legacy_label_resolves_to_original_cpsam(): + """The legacy option must still map to the original April 2025 cpsam checkpoint.""" + assert models_config.BASE_MODEL_CHECKPOINTS['cellpose-sam (legacy cpsam)'] == 'cpsam' + + +def test_default_model_is_a_selectable_base_model(): + """The configured default must be one of the offered base models.""" + assert models_config.DEFAULT_MODEL in models_config.BASE_MODELS + + +def test_base_models_offers_both_builtins(): + """Both the v2 default and the legacy option must remain available.""" + assert set(models_config.BASE_MODELS) == { + 'cellpose-sam', + 'cellpose-sam (legacy cpsam)', + } From 28d90fb6aeaec31d6c6614994b76ef8ec255caed Mon Sep 17 00:00:00 2001 From: Arjun Raj Date: Tue, 30 Jun 2026 07:19:41 -0400 Subject: [PATCH 2/2] cellposesam: guard reserved base-model names; reconcile doc convention MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the code review on this branch. - models_config.build_model_items(): new pure helper that merges the built-in labels with custom Girder model names and drops any custom name that collides with a reserved base label (compute() routes those to the built-in checkpoint, so a same-named custom model could never load). Wired into interface() in place of the inline sorted(set(...)) merge. Unit-tested (4 new cases). - CLAUDE.md: reconcile the contradictory worker-doc guidance — the "do not edit manually" bullet now matches the rest of the doc (docs are hand-maintained; the generator only stubs missing docs; the auto-doc hook is disabled; avoid --force). This is what made the CELLPOSESAM.md edit on this branch correct. The other review findings were intentional (cpsam_v2 default, ~1.2GB image growth from baking in both checkpoints) or justified (mapping pattern differs from sibling cellpose workers because label != checkpoint here), so no change. Verified: 8/8 mapping unit tests pass; worker .py files compile. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_015xrXEVpvb4c4ScEVjN1VsM --- CLAUDE.md | 7 +++- workers/annotations/cellposesam/entrypoint.py | 9 ++--- .../annotations/cellposesam/models_config.py | 14 ++++++++ .../cellposesam/tests/test_models_config.py | 35 +++++++++++++++++++ 4 files changed, 60 insertions(+), 5 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 34bd0d8..af15fae 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -47,7 +47,12 @@ When it was active, the `PostToolUse` hook (`.claude/hooks/update-worker-docs.sh ### Documentation Conventions - Each worker's doc file is named `{worker_name}.md` and lives in the worker's directory. -- Do **not** edit these files manually — they are auto-generated. Edit `entrypoint.py` instead. +- These docs are **hand-maintained**. `generate_worker_docs.py` only creates a + stub for a worker that has **no** doc file yet (existing docs are preserved, per + above), and the auto-doc hook is disabled — so edit the `{worker_name}.md` + directly when a worker's interface or behavior changes. Avoid + `generate_worker_docs.py --force`: it overwrites hand-written docs with + auto-generated stubs. - The `interfaceName`, `interfaceCategory`, `annotationShape`, and `description` Docker labels are reflected in the docs; keep these labels accurate in the Dockerfile. ## Build Commands diff --git a/workers/annotations/cellposesam/entrypoint.py b/workers/annotations/cellposesam/entrypoint.py index 26ef9ed..3439fbf 100644 --- a/workers/annotations/cellposesam/entrypoint.py +++ b/workers/annotations/cellposesam/entrypoint.py @@ -14,18 +14,19 @@ from worker_client import WorkerClient, geometry_to_polygon_coords -from models_config import BASE_MODELS, BASE_MODEL_CHECKPOINTS, DEFAULT_MODEL +from models_config import ( + BASE_MODELS, BASE_MODEL_CHECKPOINTS, DEFAULT_MODEL, build_model_items) def interface(image, apiUrl, token): client = workers.UPennContrastWorkerPreviewClient( apiUrl=apiUrl, token=token) - # models = sorted(path.stem for path in MODELS_DIR.glob('*')) - models = BASE_MODELS girder_models = [model['name'] for model in girder_utils.list_girder_models(client.client)[0]] - models = sorted(list(set(models + girder_models))) + # Reserved base labels win over any custom model of the same name (see + # build_model_items); this keeps the dropdown consistent with compute()'s routing. + models = build_model_items(girder_models) # Available types: number, text, tags, layer interface = { diff --git a/workers/annotations/cellposesam/models_config.py b/workers/annotations/cellposesam/models_config.py index df950af..1465c6e 100644 --- a/workers/annotations/cellposesam/models_config.py +++ b/workers/annotations/cellposesam/models_config.py @@ -24,3 +24,17 @@ # The list of built-in model labels offered in the interface. BASE_MODELS = list(BASE_MODEL_CHECKPOINTS) + + +def build_model_items(girder_model_names): + """Build the sorted Model-dropdown options from the built-in labels plus + custom Girder model names. + + The built-in labels are reserved: a custom Girder model whose name collides + with one is dropped, because ``compute()`` routes any name in ``BASE_MODELS`` + to the built-in checkpoint — so a same-named custom model could never be + loaded and would only create a confusing duplicate entry. + """ + custom = [name for name in girder_model_names + if name not in BASE_MODEL_CHECKPOINTS] + return sorted(set(BASE_MODELS) | set(custom)) diff --git a/workers/annotations/cellposesam/tests/test_models_config.py b/workers/annotations/cellposesam/tests/test_models_config.py index 15b17ef..6d2d08f 100644 --- a/workers/annotations/cellposesam/tests/test_models_config.py +++ b/workers/annotations/cellposesam/tests/test_models_config.py @@ -39,3 +39,38 @@ def test_base_models_offers_both_builtins(): 'cellpose-sam', 'cellpose-sam (legacy cpsam)', } + + +def test_build_model_items_includes_base_and_custom(): + """Custom Girder model names appear alongside the built-in labels.""" + items = models_config.build_model_items(['my custom model']) + assert 'cellpose-sam' in items + assert 'cellpose-sam (legacy cpsam)' in items + assert 'my custom model' in items + + +def test_build_model_items_excludes_reserved_name_collision(): + """A custom model named exactly like a base label is dropped, not duplicated. + + Otherwise it would silently route to the built-in checkpoint in compute() + and the custom weights would never be used. + """ + items = models_config.build_model_items( + ['cellpose-sam', 'cellpose-sam (legacy cpsam)']) + assert items.count('cellpose-sam') == 1 + assert items.count('cellpose-sam (legacy cpsam)') == 1 + assert set(items) == set(models_config.BASE_MODELS) + + +def test_build_model_items_sorted_and_deduped(): + """Output is sorted and free of duplicates.""" + items = models_config.build_model_items(['zeta', 'alpha', 'alpha']) + assert items == sorted(set(items)) + assert items.count('alpha') == 1 + + +def test_build_model_items_empty_returns_base_models(): + """With no custom models, only the built-in labels are offered.""" + items = models_config.build_model_items([]) + assert set(items) == set(models_config.BASE_MODELS) + assert models_config.DEFAULT_MODEL in items