Add Cellpose-SAM retrain worker#152
Conversation
Adds a cellposesam_train worker that fine-tunes Cellpose-SAM models on user-corrected annotations, mirroring the cellpose_train paradigm (training tag + optional training regions, saved to Girder for reuse by the inference worker) while adopting the Cellpose-SAM interface and training API. Key differences from the older cellpose_train worker: - Channel-slot input (Slot 1 required, Slots 2/3 optional) matching the cellposesam inference worker, instead of the cyto/nucleus Primary/Secondary channels + Nuclear Model checkbox. Cellpose-SAM ingests up to 3 channels natively, so there is no separate cytoplasm/nucleus concept. - Training via cellpose 4.x train.train_seg: no `channels` arg (uses channel_axis=-1 on channels-last stacks), AdamW (SGD deprecated), learning_rate default 1e-5, weight_decay default 0.1, bsize=256, native resolution (no diameter). - Model instantiated with CellposeModel(pretrained_model=...) from the cpsam checkpoints (models_config.py, shared with the inference worker) or a custom Girder model path; saved to .cellposesam/models. Registers the worker in build_machine_learning_workers.sh and REGISTRY.md, and includes models_config unit tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0199mMkZxJV7p7p8zTs7wYFe
cellpose 4.x train_seg saves to save_path/models/model_name and only calls (save_path/"models").mkdir(exist_ok=True) without parents=True. On the default base-model fine-tune path nothing creates /root/.cellposesam beforehand, so the save would raise FileNotFoundError. Create the models directory up front. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0199mMkZxJV7p7p8zTs7wYFe
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 80e6083b05
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if len(regionAnnotationList) == 0 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.") |
There was a problem hiding this comment.
When a user selects a Training Region tag that matches no annotations, this block warns that training will use the entire image, but it leaves training_regions non-empty. The later if training_regions is None or len(training_regions) == 0 check therefore enters the region-cropping branch, iterates an empty grouped_region_annotations[location_key], appends no samples, and calls train_seg with zero training images instead of the promised full-image fallback.
Useful? React with 👍 / 👎.
| # - 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. | ||
| model_path, train_losses, test_losses = train.train_seg( |
There was a problem hiding this comment.
Cellpose 4 train_seg defaults min_train_masks to 5, so any crop or full image with fewer than five labeled objects is removed before training. For this worker's user-corrected annotations and optional region crops, one-to-four annotated objects in a crop is a valid and likely input; those samples will be silently ignored, and if all crops are sparse the training set becomes empty. Pass an explicit lower threshold or reject sparse crops before starting training.
Useful? React with 👍 / 👎.
- Fall back to full-image training when a Training Region tag is supplied but matches no annotations. Previously the region-cropping branch was chosen based on whether a tag was given, so an unmatched tag produced zero crops and handed train_seg an empty training set despite the "using entire image" warning. The branch now hinges on whether region annotations were actually found. - Pass min_train_masks=1 to train_seg. cellpose 4's default of 5 silently drops any crop/image with fewer than five labeled objects, which is common for user-corrected crops and could empty the training set. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0199mMkZxJV7p7p8zTs7wYFe
The cellposesam_train worker inherited two latent bugs from cellpose_train; fix them at the source too: - Fall back to full-image training when a Training Region tag matches no annotations, rather than producing zero crops and training on an empty set. (Also avoids a latent len(None) crash when no region tag is set.) - Pass min_train_masks=1 so crops with fewer than five labeled objects are not silently dropped. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0199mMkZxJV7p7p8zTs7wYFe
- Remove two __pycache__/*.pyc files accidentally committed during local py_compile checks (nothing else in the repo tracks compiled artifacts). - Create .cellposesam/models once, up front in compute(), before it is used as the download destination for a custom base model and before the glob that lists it — not only just before train_seg. Avoids relying on the Girder client to create the destination directory on the custom-model path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0199mMkZxJV7p7p8zTs7wYFe
2 participants
Summary
Adds a new
cellposesam_trainannotation worker that fine-tunes Cellpose-SAM models on user-corrected annotations, producing custom models that appear directly in the Cellpose-SAM inference worker's Model dropdown.We already had a
cellpose_train(retrain for the older Cellpose) and acellposesam(inference on the latest Cellpose-SAM), but no retrain counterpart for Cellpose-SAM. This fills that gap.What's kept from the existing
cellpose_trainparadigmskimage.draw.polygon2mask(), crop to region bounding boxes, then upload the trained model to the user's Girder folder for reuse.What's new for Cellpose-SAM (differs from earlier versions)
cellposesaminference worker, replacing the cyto/nucleus Primary/Secondary channels + "Nuclear Model?" checkbox. Cellpose-SAM ingests up to three channels natively, so there is no separate cytoplasm/nucleus concept — training on the same channel layout used at inference is what matters.train.train_segAPI: nochannelsargument (stacks channels-last and passeschannel_axis=-1), AdamW optimizer (the oldSGDflag is deprecated),learning_ratedefault1e-5,weight_decaydefault0.1,bsize=256, and native-resolution training (no diameter).CellposeModel(pretrained_model=...)from thecpsam/cpsam_v2checkpoints (sharedmodels_config.py) or a custom Girder model path, replacing the removedmodel_type=argument. Models are saved to.cellposesam/models.Files
workers/annotations/cellposesam_train/:entrypoint.py,Dockerfile,environment.yml(cellpose 4.2.1.1),download_models.py,girder_utils.py,models_config.py,utils.py,CELLPOSESAM_TRAIN.md, andtests/test_models_config.py.build_machine_learning_workers.shandREGISTRY.md.Bug fixes (from code review)
Self-review + the Codex reviewer surfaced three real issues, fixed here:
train_segsaves tosave_path/models/model_namebut only calls(save_path/"models").mkdir(exist_ok=True)withoutparents=True. On the default base-model path nothing creates/root/.cellposesambeforehand (the build downloads checkpoints into cellpose's own~/.cellpose/models), so the save would raiseFileNotFoundError. The worker now creates the models directory up front.train_segan empty training set. The branch now hinges on whether region annotations were actually found, so an unmatched tag genuinely falls back to full-image training.min_train_masks. cellpose 4'strain_segdefaultsmin_train_masks=5, silently dropping any crop/image with fewer than five labeled objects — common for user-corrected crops. Now passesmin_train_masks=1.Also fixed in
cellpose_trainIssues 2 and 3 were latent in the original
cellpose_trainworker (the region-branch structure was copied from it), so the same two fixes are applied there as well.Testing
models_configmapping is unit-tested (tests/test_models_config.py) and passes in the lightweight local venv. Full Docker build/train validation on a GPU host has not been run in this environment.🤖 Generated with Claude Code
https://claude.ai/code/session_0199mMkZxJV7p7p8zTs7wYFe