cellposesam: default to cpsam_v2 checkpoint, keep legacy cpsam

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

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.

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
 

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=<name>)`` 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)

workers/annotations/cellposesam/entrypoint.py

@@ -14,18 +14,19 @@
 
 from worker_client import WorkerClient, geometry_to_polygon_coords
 
-BASE_MODELS = ['cellpose-sam']
+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 = {
@@ -68,8 +69,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 +252,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)

workers/annotations/cellposesam/environment.yml

@@ -6,5 +6,5 @@ dependencies:
 - rasterio
 - shapely
 - pip:
-  - cellpose==4.0.1
+  - cellpose==4.2.1.1
   - deeptile

workers/annotations/cellposesam/models_config.py

@@ -0,0 +1,40 @@
+"""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)
+
+
+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))

workers/annotations/cellposesam/tests/test_models_config.py

@@ -0,0 +1,76 @@
+"""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)',
+    }
+
+
+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