Add Cellpose-SAM retrain worker

REGISTRY.md

@@ -26,6 +26,7 @@ Create new annotations by segmenting images or connecting existing annotations.
 | Cellpose | This tool runs the Cellpose model to segment the image into cells. Learn more | Yes |  | [docs](workers/annotations/cellpose/CELLPOSE.md) |
 | Cellpose retrain | This tool trains a Cellpose model using user-corrected annotations. Learn more | Yes |  | [docs](workers/annotations/cellpose_train/CELLPOSE_TRAIN.md) |
 | Cellpose-SAM | This tool runs the Cellpose-SAM model to segment the image into cells. Learn more | Yes |  | [docs](workers/annotations/cellposesam/CELLPOSESAM.md) |
+| Cellpose-SAM retrain | This tool fine-tunes a Cellpose-SAM model using user-corrected annotations. Learn more | Yes | Yes | [docs](workers/annotations/cellposesam_train/CELLPOSESAM_TRAIN.md) |
 | CondensateNet segmentation | Segments biomolecular condensates using CondensateNet deep learning model | Yes |  | [docs](workers/annotations/condensatenet/CONDENSATENET.md) |
 | Connect Sequential | This tool connects objects sequentially across time or z-slices.It is useful for connec... |  | Yes | [docs](workers/annotations/connect_sequential/CONNECT_SEQUENTIAL.md) |
 | Connect Time Lapse | This tool connects objects across time slices. It allows you to connect objects even if... |  | Yes | [docs](workers/annotations/connect_timelapse/CONNECT_TIMELAPSE.md) |

build_machine_learning_workers.sh

@@ -31,6 +31,9 @@ docker build . -f ./workers/annotations/cellpose_train/Dockerfile -t annotations
 echo "Building Cellpose-SAM worker"
 docker build . -f ./workers/annotations/cellposesam/Dockerfile -t annotations/cellposesam_worker:latest $NO_CACHE
 
+echo "Building Cellpose-SAM retrain worker"
+docker build . -f ./workers/annotations/cellposesam_train/Dockerfile -t annotations/cellposesam_train_worker:latest $NO_CACHE
+
 echo "Building Stardist worker"
 docker build . -f ./workers/annotations/stardist/Dockerfile -t annotations/stardist_worker:latest $NO_CACHE
 

workers/annotations/cellpose_train/entrypoint.py

@@ -217,7 +217,14 @@ def compute(datasetId, apiUrl, token, params):
         sendError("No training annotations found.",
                   info="No annotations with the training tag were found.")
         raise ValueError("No training annotations found.")
-    if len(regionAnnotationList) == 0 and len(training_regions) > 0:
+
+    # Decide whether to crop to regions based on whether region annotations were
+    # actually found, not merely on whether a region tag was supplied. If a tag
+    # was given but matched nothing, fall back to full-image training (as the
+    # warning promises) instead of silently producing zero crops and handing
+    # train_seg an empty training set.
+    use_regions = len(regionAnnotationList) > 0
+    if not use_regions and training_regions and len(training_regions) > 0:
         sendWarning("No region annotations found.",
                     info="No annotations with the training region tag were found.")
         print("No region annotations found. Training will be performed on entire image that the annotations are in.")
@@ -267,7 +274,7 @@ def compute(datasetId, apiUrl, token, params):
             mask = draw.polygon2mask(primary_image.shape, polygon)
             label_image[mask] = i + 1
 
-        if training_regions is None or len(training_regions) == 0:
+        if not use_regions:
             training_image = np.stack(
                 [primary_image, secondary_image, np.zeros_like(primary_image)], axis=-1)
             training_images.append(training_image)
@@ -306,11 +313,15 @@ def compute(datasetId, apiUrl, token, params):
     sendProgress(0.4, "Training model",
                  f"Training with {len(training_images)} images, be patient...")
 
+    # min_train_masks defaults to 5, which would silently drop any crop/image
+    # with fewer than five labeled objects; user-corrected crops commonly have
+    # only a handful of objects, so lower it to 1 to keep every non-empty sample.
     model_path, train_losses, test_losses = train.train_seg(model.net,
                                                             train_data=training_images, train_labels=label_images,
                                                             channels=channels, normalize=True,
                                                             weight_decay=weight_decay, SGD=True, learning_rate=learning_rate,
-                                                            n_epochs=epochs, model_name=MODELS_DIR / output_model_name)
+                                                            n_epochs=epochs, min_train_masks=1,
+                                                            model_name=MODELS_DIR / output_model_name)
 
     # Upload the trained model to Girder
     sendProgress(0.95, "Saving model", f"Uploading model {output_model_name}")

workers/annotations/cellposesam_train/CELLPOSESAM_TRAIN.md

@@ -0,0 +1,103 @@
+# Cellpose-SAM Retrain Worker
+
+This worker fine-tunes a Cellpose-SAM model on user-corrected annotations, producing a custom model that can be used by the Cellpose-SAM worker.
+
+It is the Cellpose-SAM counterpart to the older [Cellpose Train worker](../cellpose_train/CELLPOSE_TRAIN.md). It follows the same annotation-gathering paradigm (training tag + optional training regions) but adopts the Cellpose-SAM interface and training API, which differ substantially from earlier Cellpose versions.
+
+## How It Works
+
+1. **Annotation Loading**: Retrieves all polygon and rectangle annotations from the dataset
+2. **Tag Filtering**: Selects training annotations by the specified training tag, and optionally crops to training region annotations
+3. **Image Assembly**: For each location (XY/Z/Time), loads the selected input-slot channels and stacks them channels-last, then renders training annotations into label masks
+4. **Region Cropping**: If training regions are specified, crops each image/mask pair to the bounding box of each region polygon
+5. **Fine-tuning**: Fine-tunes the selected base Cellpose-SAM model using `cellpose.train.train_seg()` (AdamW optimizer)
+6. **Upload**: Saves the trained model to the user's Girder `.cellposesam/models` folder
+
+## Interface Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| **Cellpose-SAM retrain** | notes | -- | Informational text with documentation link |
+| **Base Model** | select | cellpose-sam | Base model to fine-tune. `cellpose-sam` starts from the `cpsam_v2` checkpoint; `cellpose-sam (legacy cpsam)` starts from the original April 2025 `cpsam` checkpoint. User-trained models from Girder are also listed |
+| **Output Model Name** | text | -- | **Required.** Name for the saved model. Will appear in the Model dropdown of the Cellpose-SAM worker |
+| **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 |
+| **Training Tag** | tags | -- | **Required.** Tag identifying the corrected annotations to use as training ground truth |
+| **Training Region** | tags | -- | Optional tag identifying region annotations that define training crops. If empty, uses the full image |
+| **Learning Rate** | number | 0.00001 | AdamW learning rate (range: 0.000001-0.1). Cellpose-SAM fine-tunes best with a small learning rate (1e-5) |
+| **Epochs** | number | 100 | Number of training epochs (range: 10-2000) |
+| **Weight Decay** | number | 0.1 | AdamW weight decay regularization (range: 0-1) |
+
+## Implementation Details
+
+### Channel Slots vs. Primary/Secondary Channels
+
+Unlike the older `cellpose_train` worker (which used Primary/Secondary channel selectors and a "Nuclear Model?" checkbox to build the classic `channels=[cyto, nucleus]` pair), Cellpose-SAM works directly on the channels you provide. It ingests up to three channels natively and reorders/normalizes them internally, so there is no separate cytoplasm/nucleus concept.
+
+This worker therefore uses the same three `channelCheckboxes` slots as the Cellpose-SAM inference worker:
+
+- **Slot 1** is required; an error is raised if no channel is selected
+- **Slots 2 and 3** are optional
+- If multiple channels are checked in a single slot, a warning is issued and only the first is used
+- Selected channels are stacked channels-last and passed to training via `channel_axis=-1`
+
+Matching the inference worker's channel layout matters: a model trained on a given channel arrangement should be run with the same arrangement.
+
+### Training Data Preparation
+
+- Annotations are grouped by location (Time, Z, XY) to batch image loading
+- Both polygon and rectangle annotations are retrieved and combined
+- Each training annotation is rasterized into a label mask using `skimage.draw.polygon2mask()`, with each annotation assigned a unique integer label
+- Images are stacked channels-last as `(H, W, num_selected_channels)`; Cellpose-SAM standardizes to three channels internally (padding with zeros or truncating as needed)
+
+### Training Regions
+
+- Training regions are polygon/rectangle annotations with the specified region tag
+- When regions are specified, images and label masks are cropped to each region's bounding box, producing one training sample per region per location
+- When no regions are specified, the full image is used as a single training sample (a warning is displayed)
+- Using regions is recommended to focus training on relevant areas and reduce memory usage
+
+### Cellpose-SAM Training API Differences
+
+Cellpose-SAM (cellpose ≥ 4) changed the training API relative to the versions used by `cellpose_train`:
+
+- **No `channels` argument**: the network ingests channels directly, so training passes `channel_axis=-1` instead of a `channels=[...]` pair
+- **AdamW optimizer**: the old `SGD` flag is deprecated (AdamW is always used), so it is not passed
+- **Small learning rate**: the recommended default is `1e-5` (vs `0.01` for the older worker)
+- **`weight_decay=0.1`** is the Cellpose-SAM default (vs `0.0001`)
+- **`bsize=256`**: cpsam requires a block size of 256
+- **Native resolution**: training runs with `rescale=False`, so no diameter parameter is needed
+- **Model instantiation**: `CellposeModel(gpu=..., pretrained_model=<checkpoint-or-path>)` replaces the old `model_type=` argument
+
+### Model Storage
+
+- The trained model is saved locally to `/root/.cellposesam/models/` (via `save_path` + `model_name` passed to `train_seg`)
+- After training completes, the model is uploaded to the user's Girder `Private/.cellposesam/models/` folder
+- If a model with the same name already exists in Girder, it is replaced
+- Uploaded models automatically appear in the Model dropdown of the Cellpose-SAM worker
+
+### Built-in Checkpoints
+
+The base-model labels map to cellpose built-in checkpoints in `models_config.py` — `cellpose-sam` → `cpsam_v2`, `cellpose-sam (legacy cpsam)` → `cpsam`. Both checkpoints are pre-downloaded at build time by `download_models.py` so neither downloads on first run. `models_config.py` is kept in sync with the inference worker and is the single source of truth for both the interface and the build-time download.
+
+### GPU Handling
+
+The worker checks for GPU availability via `cellpose.core.use_gpu()` and instantiates the model with that result; cellpose falls back to CPU if no GPU is available.
+
+## Testing
+
+`tests/test_models_config.py` unit-tests the base-model mapping (`models_config.py`) in isolation — it has no heavy imports, so it runs in the lightweight local venv:
+
+```bash
+.cache/testvenv/bin/pytest workers/annotations/cellposesam_train/tests -q
+```
+
+## Notes
+
+- Does not use `WorkerClient` for batch processing; instead directly manages annotation retrieval and image loading (same approach as `cellpose_train`)
+- Requires a training tag and an output model name; the worker errors if either is missing
+- Training regions are optional but recommended for efficiency and to avoid training on irrelevant parts of the image
+- The base-model list is dynamically populated from both built-in checkpoints and user models in Girder
+- Progress updates are sent at key stages: loading annotations (10%), processing (20%), loading images (30%), training (40%), and saving (95%)
+- Related workers: [Cellpose-SAM](../cellposesam/CELLPOSESAM.md) (inference), [Cellpose Train](../cellpose_train/CELLPOSE_TRAIN.md) (older retrain)

workers/annotations/cellposesam_train/Dockerfile

@@ -0,0 +1,75 @@
+FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 as base
+LABEL isUPennContrastWorker=True
+LABEL com.nvidia.volumes.needed="nvidia_driver"
+
+ENV NVIDIA_VISIBLE_DEVICES all
+ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
+
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -qy tzdata && \
+  apt-get install -qy software-properties-common python3-software-properties && \
+  apt-get update && apt-get install -qy \
+  build-essential \
+  wget \
+  python3 \
+  r-base \
+  libffi-dev \
+  libssl-dev \
+  libjpeg-dev \
+  zlib1g-dev \
+  r-base \
+  git \
+  libpython3-dev && \
+  apt-get clean && rm -rf /var/lib/apt/lists/*
+
+ENV PATH="/root/miniforge3/bin:$PATH"
+ARG PATH="/root/miniforge3/bin:$PATH"
+
+RUN wget \
+    https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh \
+    && mkdir /root/.conda \
+    && bash Miniforge3-Linux-x86_64.sh -b \
+    && rm -f Miniforge3-Linux-x86_64.sh
+
+FROM base as build
+
+COPY ./workers/annotations/cellposesam_train/environment.yml /
+RUN conda env create --file /environment.yml
+SHELL ["conda", "run", "-n", "worker", "/bin/bash", "-c"]
+
+RUN pip install rtree shapely
+
+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_train/models_config.py /
+COPY ./workers/annotations/cellposesam_train/download_models.py /
+RUN python /download_models.py
+
+COPY ./workers/annotations/cellposesam_train/utils.py /
+COPY ./workers/annotations/cellposesam_train/girder_utils.py /
+COPY ./workers/annotations/cellposesam_train/entrypoint.py /
+
+COPY ./annotation_utilities /annotation_utilities
+RUN pip install /annotation_utilities
+
+COPY ./worker_client /worker_client
+RUN pip install /worker_client
+
+LABEL isUPennContrastWorker="" \
+      isAnnotationWorker="" \
+      interfaceName="Cellpose-SAM retrain" \
+      interfaceCategory="Cellpose" \
+      description="Fine-tunes Cellpose-SAM models based on user-selected images" \
+      advancedOptionsPanel="False" \
+      annotationConfigurationPanel="False" \
+      defaultToolName="Cellpose-SAM retrain" \
+      annotationShape="polygon"
+
+# Fast env activation in place of `conda run` (see todo/worker-startup-latency.md)
+COPY ./workers/base_docker_images/run_worker.sh /usr/local/bin/run_worker.sh
+RUN chmod +x /usr/local/bin/run_worker.sh
+ENTRYPOINT ["/usr/local/bin/run_worker.sh", "/entrypoint.py"]

workers/annotations/cellposesam_train/download_models.py

@@ -0,0 +1,16 @@
+"""Pre-download the built-in Cellpose-SAM checkpoints at build time.
+
+Fine-tuning starts from one of these base checkpoints, so baking them into the
+image avoids a multi-GB download on the first run. Instantiating
+``CellposeModel(pretrained_model=<name>)`` fetches the weights to
+``models.MODELS_DIR`` on first use. ``gpu=True`` matches runtime; on a GPU-less
+build host cellpose falls back to CPU but still downloads.
+"""
+
+from cellpose import models
+
+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)