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/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..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 -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) 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..1465c6e --- /dev/null +++ b/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)) 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..6d2d08f --- /dev/null +++ b/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