diff --git a/REGISTRY.md b/REGISTRY.md index 57de4e9..0cccb4e 100644 --- a/REGISTRY.md +++ b/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) | diff --git a/build_machine_learning_workers.sh b/build_machine_learning_workers.sh index 935ebdc..64c63d3 100755 --- a/build_machine_learning_workers.sh +++ b/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 diff --git a/workers/annotations/cellpose_train/entrypoint.py b/workers/annotations/cellpose_train/entrypoint.py index eb2c861..37ad229 100755 --- a/workers/annotations/cellpose_train/entrypoint.py +++ b/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}") diff --git a/workers/annotations/cellposesam_train/CELLPOSESAM_TRAIN.md b/workers/annotations/cellposesam_train/CELLPOSESAM_TRAIN.md new file mode 100644 index 0000000..e02c6fa --- /dev/null +++ b/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=)` 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) diff --git a/workers/annotations/cellposesam_train/Dockerfile b/workers/annotations/cellposesam_train/Dockerfile new file mode 100644 index 0000000..8fd873e --- /dev/null +++ b/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"] diff --git a/workers/annotations/cellposesam_train/download_models.py b/workers/annotations/cellposesam_train/download_models.py new file mode 100644 index 0000000..6c0bfd6 --- /dev/null +++ b/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=)`` 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) diff --git a/workers/annotations/cellposesam_train/entrypoint.py b/workers/annotations/cellposesam_train/entrypoint.py new file mode 100644 index 0000000..322613b --- /dev/null +++ b/workers/annotations/cellposesam_train/entrypoint.py @@ -0,0 +1,400 @@ +import argparse +from collections import defaultdict +import json +import sys + +from skimage import draw +import numpy as np + +from shapely.geometry import Polygon + +import annotation_client.workers as workers +import annotation_client.tiles as tiles +import annotation_client.annotations as annotations +from annotation_client.utils import sendProgress, sendError, sendWarning +import annotation_utilities.annotation_tools as annotation_tools + +import girder_utils +from girder_utils import CELLPOSE_DIR, MODELS_DIR + +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) + + girder_models = [model['name'] + for model in girder_utils.list_girder_models(client.client)[0]] + # 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 = { + 'Cellpose-SAM retrain': { + 'type': 'notes', + 'value': 'This tool fine-tunes a Cellpose-SAM model using user-corrected annotations. ' + 'Cellpose-SAM works directly on the channels you provide (up to three), so there ' + 'is no separate cytoplasm/nucleus channel concept — just pick the channels you ' + 'want the model to see. ' + 'Learn more', + 'displayOrder': 0, + }, + 'Base Model': { + 'type': 'select', + 'items': models, + 'default': DEFAULT_MODEL, + 'tooltip': 'The model used as a starting point for fine-tuning.\n' + '"cellpose-sam" starts from the current default checkpoint (cpsam_v2).\n' + '"cellpose-sam (legacy cpsam)" starts from the original April 2025 model.\n' + 'Custom models you have previously trained are also listed here.', + 'noCache': True, + 'displayOrder': 1, + }, + 'Output Model Name': { + 'type': 'text', + 'tooltip': 'The name of the retrained model (saved to your .cellposesam/models folder).\n' + 'It will appear in the Model dropdown of the Cellpose-SAM worker.', + 'displayOrder': 2, + }, + 'Channel for Slot 1': { + 'type': 'channelCheckboxes', + 'tooltip': "Select source channel(s) for the model's first input slot. If multiple are selected, only the first will be used. This slot is required.", + 'displayOrder': 3, + }, + 'Channel for Slot 2': { + 'type': 'channelCheckboxes', + 'tooltip': "Select source channel(s) for the model's second input slot. If multiple are selected, only the first will be used. (Optional)", + 'displayOrder': 4, + }, + 'Channel for Slot 3': { + 'type': 'channelCheckboxes', + 'tooltip': "Select source channel(s) for the model's third input slot. If multiple are selected, only the first will be used. (Optional)", + 'displayOrder': 5, + }, + 'Training Tag': { + 'type': 'tags', + 'tooltip': 'Train the model on objects that have this tag.', + 'displayOrder': 6, + }, + 'Training Region': { + 'type': 'tags', + 'tooltip': 'These objects define the regions that the training will be performed on.\n' + 'If you do not select any objects, the training will be performed on all objects in the image.\n' + 'You can and probably should select multiple regions.', + 'displayOrder': 7, + }, + 'Learning Rate': { + 'type': 'number', + 'min': 0.000001, + 'max': 0.1, + 'default': 0.00001, + 'tooltip': 'The learning rate for fine-tuning. Cellpose-SAM fine-tunes best with a small ' + 'learning rate; 1e-5 (0.00001) is the recommended default.', + 'displayOrder': 8, + }, + 'Epochs': { + 'type': 'number', + 'min': 10, + 'max': 2000, + 'default': 100, + 'tooltip': 'The number of epochs to train the model for. 100 is a good starting point for Cellpose-SAM.', + 'displayOrder': 9, + }, + 'Weight Decay': { + 'type': 'number', + 'min': 0, + 'max': 1, + 'default': 0.1, + 'tooltip': 'The weight decay (AdamW regularization) for training. 0.1 is the Cellpose-SAM default.', + 'displayOrder': 10, + }, + } + # Send the interface object to the server + client.setWorkerImageInterface(image, interface) + + +def get_slot_channels(workerInterface): + """Resolve the selected input-slot channels into an ordered channel list. + + Mirrors the cellposesam inference worker: Slot 1 is required, Slots 2 and 3 + are optional, and if multiple channels are checked in a slot only the first + is used (with a warning). + """ + slot1 = [k for k, v in workerInterface.get('Channel for Slot 1', {}).items() if v] + slot2 = [k for k, v in workerInterface.get('Channel for Slot 2', {}).items() if v] + slot3 = [k for k, v in workerInterface.get('Channel for Slot 3', {}).items() if v] + + stack_channels = [] + + if not slot1: + sendError("No channel selected for Slot 1. This is a required field.") + raise ValueError("No channel selected for Slot 1.") + if len(slot1) > 1: + sendWarning( + f"Multiple channels selected for Slot 1 ({slot1}). Using the first: {slot1[0]}.") + stack_channels.append(int(slot1[0])) + + if slot2: + if len(slot2) > 1: + sendWarning( + f"Multiple channels selected for Slot 2 ({slot2}). Using the first: {slot2[0]}.") + stack_channels.append(int(slot2[0])) + + if slot3: + if len(slot3) > 1: + sendWarning( + f"Multiple channels selected for Slot 3 ({slot3}). Using the first: {slot3[0]}.") + stack_channels.append(int(slot3[0])) + + return stack_channels + + +def compute(datasetId, apiUrl, token, params): + """ + params (could change): + configurationId, + datasetId, + description: tool description, + type: tool type, + id: tool id, + name: tool name, + image: docker image, + channel: annotation channel, + assignment: annotation assignment ({XY, Z, Time}), + tags: annotation tags (list of strings), + tile: tile position (TODO: roi) ({XY, Z, Time}), + connectTo: how new annotations should be connected + """ + + # Lazy import: keeps cellpose off the interface/startup path (~seconds). See todo/worker-startup-latency.md + from cellpose import models, train, core + + workerInterface = params['workerInterface'] + + # Get the model and training parameters from interface values + base_model = workerInterface['Base Model'] + output_model_name = workerInterface['Output Model Name'] + training_tag = workerInterface.get('Training Tag', None) + training_regions = workerInterface.get('Training Region', None) + learning_rate = float(workerInterface['Learning Rate']) + epochs = int(workerInterface['Epochs']) + weight_decay = float(workerInterface['Weight Decay']) + + print(f"Training tag: {training_tag}") + print(f"Training regions: {training_regions}") + + if not output_model_name: + sendError("No output model name provided.", + info="Please provide a name for the retrained model.") + raise ValueError("No output model name provided.") + + if training_tag is None or len(training_tag) == 0: + sendError("No training tag selected.", + info="Choose a tag for training annotations.") + raise ValueError("No training tag selected.") + if training_regions is None or len(training_regions) == 0: + sendWarning("No training regions selected.", + info="Training will be performed on entire image that the annotations are in.") + print("No training regions selected. Training will be performed on entire image that the annotations are in.") + + # Resolve the input-slot channels (Slot 1 required, Slots 2 & 3 optional). + stack_channels = get_slot_channels(workerInterface) + print(f"Using channels for Cellpose-SAM input (slots 1, 2, 3): {stack_channels}") + + client = workers.UPennContrastWorkerPreviewClient( + apiUrl=apiUrl, token=token) + annotationClient = annotations.UPennContrastAnnotationClient( + apiUrl=apiUrl, token=token) + tileClient = tiles.UPennContrastDataset( + apiUrl=apiUrl, token=token, datasetId=datasetId) + + # Ensure the local models directory exists up front. It is used both as the + # download destination for a custom base model (below) and as the training + # save location: train_seg saves to `save_path / "models" / model_name` and + # only does `(save_path / "models").mkdir(exist_ok=True)` (no parents), so the + # parent `.cellposesam` dir must already exist. Unlike the older cellpose_train + # worker (which saved into cellpose's own ~/.cellpose/models), our custom + # .cellposesam dir is never created during the build or on the base-model path. + MODELS_DIR.mkdir(parents=True, exist_ok=True) + + # Resolve the base checkpoint / model path to fine-tune from. + if base_model in BASE_MODELS: + # Pass the checkpoint name explicitly so the starting point is pinned to + # the selected model rather than cellpose's internal default, which can + # change between versions. + pretrained_model = BASE_MODEL_CHECKPOINTS[base_model] + else: + girder_utils.download_girder_model(client.client, base_model) + pretrained_model = str(MODELS_DIR / base_model) + + # Print the contents of the models directory + print(f"Models directory contents: {list(MODELS_DIR.glob('*'))}") + + # Initial loading phase + sendProgress(0.1, "Loading annotations", + "Retrieving annotations from server") + blobAnnotationList = annotationClient.getAnnotationsByDatasetId( + datasetId, limit=1000000, shape='polygon') + rectangleAnnotationList = annotationClient.getAnnotationsByDatasetId( + datasetId, limit=1000000, shape='rectangle') + # Add the rectangle annotations to the blob annotations + blobAnnotationList.extend(rectangleAnnotationList) + + trainingAnnotationList = annotation_tools.get_annotations_with_tags( + blobAnnotationList, training_tag, exclusive=False) + if training_regions is None or len(training_regions) == 0: + regionAnnotationList = [] + else: + regionAnnotationList = annotation_tools.get_annotations_with_tags( + blobAnnotationList, training_regions, exclusive=False) + print(f"Training on {len(regionAnnotationList)} region annotations.") + + if len(trainingAnnotationList) == 0: + sendError("No training annotations found.", + info="No annotations with the training tag were found.") + raise ValueError("No training annotations found.") + + # 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.") + + sendProgress(0.2, "Processing annotations", + "Grouping annotations by location") + # Group the training annotations by location so that we can batch the image loading. + grouped_training_annotations = defaultdict(list) + for current_annotation in trainingAnnotationList: + location_key = (current_annotation['location']['Time'], + current_annotation['location']['Z'], current_annotation['location']['XY']) + grouped_training_annotations[location_key].append(current_annotation) + + # Group the region annotations by location so that we can batch the image loading. + grouped_region_annotations = defaultdict(list) + for current_annotation in regionAnnotationList: + location_key = (current_annotation['location']['Time'], + current_annotation['location']['Z'], current_annotation['location']['XY']) + grouped_region_annotations[location_key].append(current_annotation) + + training_images = [] + label_images = [] + + sendProgress(0.3, "Loading training data", "Loading training images") + # Loop through each location and load the image for the training. + for location_key, training_annotations in grouped_training_annotations.items(): + time, z, xy = location_key + + # Load and stack the selected channels (channels-last) for this location. + channel_images = [] + for channel in stack_channels: + frame = tileClient.coordinatesToFrameIndex(xy, z, time, channel) + channel_image = tileClient.getRegion(datasetId, frame=frame).squeeze() + channel_images.append(channel_image) + # Shape: (H, W, num_selected_channels). Cellpose-SAM standardizes to + # three channels internally (padding with zeros / truncating as needed). + stacked_image = np.stack(channel_images, axis=-1) + + label_image = np.zeros(stacked_image.shape[:2], dtype=np.uint16) + for i, current_annotation in enumerate(training_annotations): + polygon = np.array([list(coordinate.values())[1::-1] + for coordinate in current_annotation['coordinates']]) + mask = draw.polygon2mask(label_image.shape, polygon) + label_image[mask] = i + 1 + + if not use_regions: + training_images.append(stacked_image) + label_images.append(label_image) + else: + region_annotations = grouped_region_annotations[location_key] + for region_annotation in region_annotations: + region_polygon = Polygon([(coordinate['x'], coordinate['y']) + for coordinate in region_annotation['coordinates']]) + + # Use shapely to get the bounding box + min_x, min_y, max_x, max_y = region_polygon.bounds + + # Crop both the stacked image and the label mask to the bounding box. + stacked_image_crop = stacked_image[int( + min_y):int(max_y), int(min_x):int(max_x)] + label_image_crop = label_image[int(min_y):int( + max_y), int(min_x):int(max_x)] + + training_images.append(stacked_image_crop) + label_images.append(label_image_crop) + + using_gpu = core.use_gpu() + print(f"Using GPU: {using_gpu}") + + # Cellpose-SAM (cellpose >= 4) drops model_type in favor of pretrained_model; + # passing 'cpsam'/'cpsam_v2' or a custom model path starts fine-tuning from + # those weights. + model = models.CellposeModel(gpu=using_gpu, pretrained_model=pretrained_model) + + print(f"Training with {len(training_images)} images.") + sendProgress(0.4, "Training model", + f"Training with {len(training_images)} images, be patient...") + + # Cellpose-SAM training differs from earlier Cellpose versions: + # - no `channels` argument (the network ingests up to 3 channels directly); + # channel_axis=-1 tells train_seg our stacks are channels-last. + # - AdamW is always used (the old SGD flag is deprecated), so we do not pass SGD. + # - bsize must be 256 for cpsam. + # - training uses native resolution (rescale=False), so no diameter is needed. + # - 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 (samples with zero labels are still excluded by cellpose). + model_path, train_losses, test_losses = train.train_seg( + model.net, + train_data=training_images, + train_labels=label_images, + channel_axis=-1, + normalize=True, + weight_decay=weight_decay, + learning_rate=learning_rate, + n_epochs=epochs, + bsize=256, + min_train_masks=1, + save_path=str(CELLPOSE_DIR), + model_name=output_model_name) + + print(f"Trained model saved to: {model_path}") + + # Upload the trained model to Girder + sendProgress(0.95, "Saving model", f"Uploading model {output_model_name}") + girder_utils.upload_girder_model(client.client, output_model_name) + + +if __name__ == '__main__': + # Define the command-line interface for the entry point + parser = argparse.ArgumentParser( + description='Fine-tune a Cellpose-SAM model on user-corrected annotations') + + parser.add_argument('--datasetId', type=str, + required=False, action='store') + parser.add_argument('--apiUrl', type=str, required=True, action='store') + parser.add_argument('--token', type=str, required=True, action='store') + parser.add_argument('--request', type=str, required=True, action='store') + parser.add_argument('--parameters', type=str, + required=True, action='store') + + args = parser.parse_args(sys.argv[1:]) + + params = json.loads(args.parameters) + datasetId = args.datasetId + apiUrl = args.apiUrl + token = args.token + + match args.request: + case 'compute': + compute(datasetId, apiUrl, token, params) + case 'interface': + interface(params['image'], apiUrl, token) diff --git a/workers/annotations/cellposesam_train/environment.yml b/workers/annotations/cellposesam_train/environment.yml new file mode 100644 index 0000000..e80333c --- /dev/null +++ b/workers/annotations/cellposesam_train/environment.yml @@ -0,0 +1,10 @@ +name: worker +dependencies: +- python=3.10 +- pip +- imageio +- rasterio +- shapely +- pip: + - cellpose==4.2.1.1 + - scikit-image diff --git a/workers/annotations/cellposesam_train/girder_utils.py b/workers/annotations/cellposesam_train/girder_utils.py new file mode 100644 index 0000000..6ce4f54 --- /dev/null +++ b/workers/annotations/cellposesam_train/girder_utils.py @@ -0,0 +1,66 @@ +from pathlib import Path + +# Assuming these are defined elsewhere in your cellpose worker +CELLPOSE_DIR = Path('/root/.cellposesam') +MODELS_DIR = CELLPOSE_DIR / 'models' + + +def mkdir(gc, parent_id, folder_name): + + folders = gc.get('folder', parameters={'parentId': parent_id, + 'parentType': 'folder', 'name': folder_name}) + + if folders: + _id = folders[0]['_id'] + else: + new_folder = gc.post('folder', parameters={ + 'parentId': parent_id, + 'parentType': 'folder', + 'name': folder_name + }) + _id = new_folder['_id'] + + return _id + + +def get_cellpose_dir(gc): + + user_id = gc.get('user/me')['_id'] + private_folder_id = gc.get('folder', parameters={ + 'parentId': user_id, 'parentType': 'user', 'name': 'Private'})[0]['_id'] + cellpose_folder_id = mkdir(gc, private_folder_id, '.cellposesam') + + return cellpose_folder_id + + +def list_girder_models(gc): + + cellpose_folder_id = get_cellpose_dir(gc) + models_folder_id = mkdir(gc, cellpose_folder_id, 'models') + girder_models = list(gc.listItem(models_folder_id)) + + return girder_models, models_folder_id + + +def list_local_models(): + """List models in local .cellposesam/models directory""" + MODELS_DIR.mkdir(parents=True, exist_ok=True) + return list(MODELS_DIR.glob('*')) + + +def download_girder_model(gc, model_name): + + girder_models, _ = list_girder_models(gc) + girder_model = [model for model in girder_models if model['name'] == model_name] + if girder_model: + gc.downloadItem(girder_model[0]['_id'], MODELS_DIR, girder_model[0]['name']) + + +def upload_girder_model(gc, model_name): + + girder_models, models_folder_id = list_girder_models(gc) + girder_model = [model for model in girder_models if model['name'] == model_name] + if girder_model: + gc.delete(f"{girder_model[0]['_modelType']}/{girder_model[0]['_id']}") + + gc.uploadFileToFolder(models_folder_id, MODELS_DIR / model_name) diff --git a/workers/annotations/cellposesam_train/models_config.py b/workers/annotations/cellposesam_train/models_config.py new file mode 100644 index 0000000..9c9dc0f --- /dev/null +++ b/workers/annotations/cellposesam_train/models_config.py @@ -0,0 +1,42 @@ +"""Built-in Cellpose-SAM model options for the cellposesam_train worker. + +Maps the human-friendly names shown in the Base Model dropdown to the cellpose +checkpoint identifiers passed to ``CellposeModel(pretrained_model=...)`` as the +starting point for fine-tuning. + +Kept in sync with the cellposesam (inference) worker's ``models_config.py`` so +the base-model choices match. Kept deliberately import-free (no +cellpose/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 Base 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_train/tests/__init__.py b/workers/annotations/cellposesam_train/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/workers/annotations/cellposesam_train/tests/test_models_config.py b/workers/annotations/cellposesam_train/tests/test_models_config.py new file mode 100644 index 0000000..16dfae4 --- /dev/null +++ b/workers/annotations/cellposesam_train/tests/test_models_config.py @@ -0,0 +1,76 @@ +"""Unit tests for the built-in Cellpose-SAM base-model mapping. + +These exercise ``models_config`` in isolation — it must stay free of heavy +imports (cellpose/annotation_client) so it runs in the lightweight local venv +without the full worker stack. Run with: + + .cache/testvenv/bin/pytest workers/annotations/cellposesam_train/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 fine-tune from the 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 as the fine-tuning base. + """ + 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 diff --git a/workers/annotations/cellposesam_train/utils.py b/workers/annotations/cellposesam_train/utils.py new file mode 100644 index 0000000..31b5e43 --- /dev/null +++ b/workers/annotations/cellposesam_train/utils.py @@ -0,0 +1,50 @@ +from itertools import chain + + +def process_range_list(rl): + + g = parse_range_list(rl) + first, g = peek_generator(g) + if first is None: + g = None + + return g + + +def parse_range_list(rl): + ranges = sorted(set(map(_parse_range, rl.split(','))), key=lambda x: (x.start, x.stop)) + return chain.from_iterable(_collapse_range(ranges)) + + +def peek_generator(g): + + first = next(g, None) + g = chain([first], g) + + return first, g + + +def _parse_range(r): + parts = list(_split_range(r.strip())) + if len(parts) == 0: + return range(0, 0) + elif len(parts) > 2: + raise ValueError('Invalid range: {}'.format(r)) + return range(parts[0], parts[-1] + 1) + + +def _collapse_range(ranges): + end = None + for value in ranges: + yield range(max(end, value.start), max(value.stop, end)) if end else value + end = max(end, value.stop) if end else value.stop + + +def _split_range(value): + value = value.split('-') + for val, prev in zip(value, chain((None,), value)): + if val != '': + val = int(val) + if prev == '': + val *= -1 + yield val