Skip to content

Commit 68a2123

Browse files
danifrancodanifranco
committed
Add same3 encoder example
1 parent 67a780d commit 68a2123

6 files changed

Lines changed: 763 additions & 6 deletions

File tree

‎example/export_cellpose_model/README.md‎

Lines changed: 185 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,190 @@
1-
Here are the scripts used to create a bioimage.io export of the Cellpose-SAM model.
1+
# Cellpose-SAM
22

3-
1. `cellpose_original.py` Run original cellpose model and save an analog input and output for bioimageio tests
4-
1. `bioimageio_export` Create a bioimage.io model description, then export and test it.
5-
1. `analyze_export.py` Compare cellpose results to bioimage.io results in depth.
3+
A generalist deep-learning model for 2D cellular instance segmentation. Cellpose-SAM adapts the pretrained Vision Transformer backbone of the Segment Anything Model (SAM) into the Cellpose framework, producing a foundation model that segments cells and nuclei across a very broad range of imaging modalities, cell types and organisms without retraining.
4+
5+
This repository packages the original Cellpose-SAM weights (`cpsam`) as a [BioImage Model Zoo](https://bioimage.io) model so it can be run directly from BioImage.IO-compatible software.
6+
7+
# Table of Contents
8+
9+
- [Model Details](#model-details)
10+
- [Uses](#uses)
11+
- [Task Details](#task-details)
12+
- [Bias, Risks, and Limitations](#bias-risks-and-limitations)
13+
- [Training Details](#training-details)
14+
- [Evaluation](#evaluation)
15+
- [Technical Specifications](#technical-specifications)
16+
- [How to Get Started with the Model](#how-to-get-started-with-the-model)
17+
- [Reproducing this BioImage.IO export](#reproducing-this-bioimageio-export)
18+
19+
# Model Details
20+
21+
## Model Description
22+
23+
Cellpose-SAM replaces the convolutional backbone used in earlier Cellpose releases with the ViT-L transformer encoder from Meta's Segment Anything Model, fine-tuned inside the Cellpose spatial-flow segmentation framework. The network predicts spatial flow fields and a cell-probability map; these are converted into a labelled instance mask by the Cellpose flow-dynamics post-processing. The authors report that the model substantially exceeds inter-human annotation agreement and approaches the human-consensus bound, and that it is explicitly trained to be robust to channel shuffling, cell size, shot noise, downsampling, and isotropic/anisotropic blur.
24+
25+
- **Developed by:** Marius Pachitariu, Michael Rariden and Carsen Stringer (Stringer & Pachitariu labs, HHMI Janelia Research Campus)
26+
- **Shared by:** Packaged for the BioImage Model Zoo by Daniel Franco-Barranco
27+
- **Model type:** Vision Transformer (SAM ViT-L backbone) within the Cellpose segmentation framework
28+
- **Modality:** 2D light microscopy — fluorescence, brightfield, phase contrast, and other cellular imaging
29+
- **License:** [BSD-3-Clause](https://github.com/MouseLand/cellpose/blob/main/LICENSE) (Cellpose); SAM backbone originally released under Apache-2.0
30+
- **Finetuned from model:** Segment Anything Model (SAM), ViT-L image encoder
31+
32+
## Model Sources
33+
34+
- **Repository:** https://github.com/MouseLand/cellpose
35+
- **Paper:** Pachitariu, Rariden & Stringer, *Cellpose-SAM: superhuman generalization for cellular segmentation*, bioRxiv 2025. DOI: [10.1101/2025.04.28.651001](https://doi.org/10.1101/2025.04.28.651001)
36+
- **Weights:** https://huggingface.co/mouseland/cellpose-sam
37+
- **Documentation:** https://cellpose.readthedocs.io
38+
39+
# Uses
40+
41+
## Direct Use
42+
43+
The model is a **generalist segmenter** intended to be used out of the box, without fine-tuning, on 2D microscopy images of cells or nuclei. It accepts up to three input channels (e.g. cytoplasm + nucleus, RGB histology, or a single grayscale channel replicated across channels) and returns an instance-label image where `0` is background and each connected object receives a unique integer ID.
44+
45+
Because it was trained with heavy channel and appearance augmentation, it generalizes across many microscope types, stains and organisms without the user needing to specify a "cytoplasm" vs "nucleus" channel order.
46+
47+
## Downstream Use
48+
49+
Cellpose-SAM fits into the wider Cellpose ecosystem and can be:
50+
51+
- fine-tuned or trained human-in-the-loop on domain-specific data (see the Cellpose GUI and `cellpose` training API),
52+
- combined with Cellpose image-restoration models,
53+
- used as the 2D engine for 3D segmentation via slice-wise stitching.
54+
55+
## Out-of-Scope Use
56+
57+
- **Not for clinical or diagnostic use.** The model is a research tool and has not been validated for medical decision-making.
58+
- Not validated for imaging modalities far outside the training distribution (e.g. raw electron microscopy volumes) — validate before relying on it.
59+
- Native inference here is **2D**; true 3D volumetric segmentation requires the Cellpose 3D pipeline rather than this packaged tile-based model.
60+
- Always verify results on your own data before drawing quantitative conclusions.
61+
62+
# Task Details
63+
64+
- **Task type:** Instance segmentation
65+
- **Input modality:** 2D light microscopy (fluorescence, brightfield, phase contrast, histology)
66+
- **Target structures:** Cells and nuclei
67+
- **Imaging technique:** Broadly modality-agnostic by design
68+
- **Spatial resolution:** No fixed pixel size; the model is trained to be robust to cell size and to down-sampling. Very small or very large objects relative to the tile may benefit from rescaling.
69+
70+
# Bias, Risks, and Limitations
71+
72+
## Known Biases
73+
74+
- Training data, while diverse, is dominated by 2D light-microscopy of common model systems; rare or highly unusual morphologies may be under-represented.
75+
- Performance still depends on image quality; extreme noise, artefacts or out-of-focus regions can degrade results.
76+
77+
## Limitations
78+
79+
- Objects that are heavily touching, overlapping or transparent remain challenging.
80+
- The packaged model runs on fixed-size tiles (256×256) with a 16-pixel halo; very large images are processed tile-by-tile and objects spanning tile borders rely on the halo for continuity.
81+
- Results depend on the flow-dynamics thresholds (`flow_threshold`, `cellprob_threshold`, `min_size`), which are fixed in this package but can be tuned in the native Cellpose API.
82+
83+
## Recommendations
84+
85+
- Validate on a representative subset of your own images.
86+
- Where accuracy is critical, use manual verification or human-in-the-loop correction.
87+
- For non-standard data, consider fine-tuning within Cellpose rather than relying on the generalist weights alone.
88+
89+
# Training Details
90+
91+
Full training methodology, datasets and hyperparameters are described in the original paper and repository. In brief:
92+
93+
- **Backbone:** SAM ViT-L image encoder, adapted into the Cellpose flow-field framework.
94+
- **Augmentations (key contribution):** channel shuffling, cell-size variation, shot noise, down-sampling, and isotropic/anisotropic blur, to maximize generalization.
95+
- **Objective:** Cellpose spatial-flow + cell-probability loss.
96+
- **Framework:** PyTorch.
97+
98+
For exact dataset composition, splits, epochs and optimizer settings, refer to the [paper](https://doi.org/10.1101/2025.04.28.651001) and the [Cellpose repository](https://github.com/MouseLand/cellpose).
99+
100+
# Evaluation
101+
102+
The original authors evaluate Cellpose-SAM against inter-human annotation agreement and report that it **exceeds inter-human agreement and approaches the human-consensus segmentation bound** across a diverse benchmark of cellular images, outperforming previous Cellpose generations. Quantitative benchmarks (Average Precision at IoU thresholds and comparisons to prior models) are reported in the paper.
103+
104+
This BioImage.IO package is additionally checked for numerical reproducibility: the exported model's label output is compared against the reference Cellpose inference on the bundled sample image, within the tolerance declared in the model description (see `bioimageio_export.py`).
105+
106+
# Technical Specifications
107+
108+
## Model Architecture and Objective
109+
110+
- **Architecture:** SAM ViT-L transformer encoder within the Cellpose segmentation head (`CPSAM`); loaded here through the BioImage.IO-compatible subclass `CPnetBioImageIO`.
111+
- **Objective:** Predicts spatial flow fields + cell-probability map, decoded into instances by Cellpose flow dynamics.
112+
113+
### Input specification
114+
115+
| Property | Value |
116+
|----------|-------|
117+
| Axes | `batch`, `channel`, `y`, `x` |
118+
| Channels | 3 (`channel0`, `channel1`, `channel2`) |
119+
| Tile size | 256 × 256 |
120+
| Data type | float32 |
121+
| Padding | constant, value `0` |
122+
| Preprocessing | percentile scale-range normalization (1st–99th percentile over `batch, y, x`, `eps=1e-8`) |
123+
124+
### Output specification
125+
126+
| Property | Value |
127+
|----------|-------|
128+
| Tensor | `labels` — instance label image, `0` = background |
129+
| Axes | `batch`, `channel`, `y`, `x` (16-pixel halo on `y` and `x`) |
130+
| Postprocessing | Cellpose flow dynamics (`cellprob_threshold=0.0`, `flow_threshold=0.4`, `do_3D=False`, `min_size=15`) |
131+
132+
## Compute Infrastructure
133+
134+
### Hardware Requirements
135+
136+
- **Inference:** Runs on CPU; a CUDA GPU is strongly recommended for reasonable speed on large images.
137+
- **Storage:** The `cpsam` weights are several hundred MB.
138+
139+
### Software Dependencies
140+
141+
- **Framework:** PyTorch (tested with 2.10.0)
142+
- **Libraries:** `cellpose==4.2.1.1` (pinned in the packaged `environment.yaml`), `bioimageio.core`
143+
- **BioImage.IO compatibility:** loadable via `bioimageio.core`; the weights are provided as a PyTorch state dict with a bundled architecture file, so no upstream Cellpose patch is required.
144+
145+
# How to Get Started with the Model
146+
147+
Using the BioImage.IO Python API:
148+
149+
```python
150+
from bioimageio.core import load_description, predict
151+
import numpy as np
152+
153+
model = load_description("cpsam_bioimageio.zip") # or the bioimage.io model id
154+
155+
# image: (batch, channel=3, y, x) float array; single-channel data can be
156+
# broadcast across the three channels.
157+
image = np.load("test_input.npy")
158+
labels = predict(model=model, inputs={"input": image})
159+
```
160+
161+
To run the model natively (full Cellpose feature set, tuning, 3D, training), install Cellpose and follow the [official documentation](https://cellpose.readthedocs.io):
162+
163+
```bash
164+
pip install "cellpose==4.2.1.1"
165+
```
166+
167+
# Reproducing this BioImage.IO export
168+
169+
The scripts in this folder build and validate the BioImage.IO package from the original Cellpose-SAM weights:
170+
171+
1. `cellpose_original.py` — runs the original Cellpose model and saves matching input/output tensors used as BioImage.IO test tensors.
172+
2. `bioimageio_export.py` — builds the `ModelDescr`, then exports and tests the packaged model.
173+
3. `analyze_export.py` — compares Cellpose results against the BioImage.IO results in depth.
174+
4. `cellpose_vit_bioimageio.py` — the bundled `CPnetBioImageIO` architecture that lets the exported model run against a stock, unmodified Cellpose install.
6175

7176
```console
8177
python cellpose_original.py
9178
python bioimageio_export.py
10-
python analzse_export.py
11-
````
179+
python analyze_export.py
180+
```
181+
182+
# Citation
183+
184+
If you use this model, please cite the original work:
185+
186+
> Pachitariu, M., Rariden, M., & Stringer, C. (2025). *Cellpose-SAM: superhuman generalization for cellular segmentation.* bioRxiv. https://doi.org/10.1101/2025.04.28.651001
187+
188+
# License
189+
190+
This model is distributed under the [BSD-3-Clause license](https://github.com/MouseLand/cellpose/blob/main/LICENSE) of the Cellpose project. The SAM backbone from which it is fine-tuned was originally released under Apache-2.0.

‎example/export_sam3_encoder/README.md‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
# SAM 3 image encoder → bioimage.io
2+
3+
Scripts to export the **SAM 3 image encoder** — the vision backbone of Meta's
4+
["Segment Anything with Concepts"](https://github.com/facebookresearch/sam3)
5+
(SAM 3, Nov 2025) — as a [bioimage.io](https://bioimage.io) model. The exported
6+
model maps a `1008×1008` RGB image to the backbone `vision_features` map, shape
7+
`(1, 256, 72, 72)`.
8+
9+
Modeled after [`../export_cellpose_model`](../export_cellpose_model).
10+
11+
## Setup
12+
13+
SAM 3 is **not** bundled with this example — clone and install it yourself
14+
([facebookresearch/sam3](https://github.com/facebookresearch/sam3)):
15+
16+
```bash
17+
git clone https://github.com/facebookresearch/sam3.git
18+
pip install ./sam3
19+
```
20+
21+
The checkpoint (`sam3.pt`, ~3.45 GB) is the gated HuggingFace `facebook/sam3`
22+
weight; run `hf auth login` and download it first (see `sam3_original.py` for the
23+
default cache path).
24+
25+
## Environment
26+
27+
`sam3` and `bioimageio.core` must be importable **together**, on Python ≥ 3.11 with a CUDA GPU. `bioimageio_export.py` regenerates `output/environment.yaml` with the full pinned list; the non-obvious pieces, all discovered the hard way:
28+
29+
* **`numpy<2`** — SAM 3 requires it; bioimage.io is happy with `numpy>=1.26`, so
30+
the two coexist.
31+
* **`setuptools<81`** — SAM 3 still imports `pkg_resources`, which setuptools
32+
removed in v81.
33+
* **SAM 3's load-time extras** — `sam3.model_builder` imports several `[train]`
34+
`[notebooks]` extras at import time (`einops`, `submitit`, `tensorboard`,
35+
`fvcore`, `fairscale`, `torchmetrics`, `zstandard`, `hydra-core`, `psutil`,
36+
`pycocotools`, …). A plain `pip install git+…sam3` pulls none of them.
37+
38+
The visualization / video / eval / fast-kernel deps (`cv2`, `decord`,
39+
`matplotlib`, `xformers`, `flash_attn`, `detectron2`, …) are **not** needed — the
40+
image encoder never imports them on its path.
41+
42+
## How it works / key design decisions
43+
44+
* **Encoder-only wrapper.** `build_sam3_image_model(..., enable_segmentation=
45+
False, enable_inst_interactivity=False)` builds the full image model without
46+
weights; the wrapper keeps only `.backbone` and returns its top-level
47+
`vision_features`. `load_state_dict` strips the checkpoint's `detector.` prefix
48+
and keeps the `backbone.*` tensors (mirroring SAM 3's own `_load_checkpoint`).
49+
* **BFloat16 on GPU.** SAM 3's fused `addmm_act` kernel forces BFloat16 (and, in
50+
practice, CUDA). The wrapper holds the backbone in BFloat16 so weights match
51+
activations; bioimage.io loads the float32 checkpoint into those BFloat16
52+
params, casting on copy — the same rounding used to make the reference tensors.
53+
* **Normalisation lives in the wrapper, not in bioimage.io preprocessing.** SAM
54+
3's BFloat16 stack is chaotically input-sensitive: a ~1e-5 gap (e.g. float32
55+
rounding between numpy and bioimage.io's `scale_linear`+`normalize`) blows up
56+
into ~1.0 output differences over ~1% of elements — more than the bioimage.io
57+
reproducibility-tolerance caps can absorb. So the model declares **no**
58+
preprocessing; `forward` takes a raw `[0, 255]` RGB tensor and does the
59+
scale-to-[0,1] + ImageNet mean/std itself. The identical torch code then
60+
produces the model input both when generating the reference tensors and when
61+
bioimage.io validates them, so there is nothing to amplify.
62+
* **Contiguous input.** The BFloat16 kernels are also layout sensitive (a
63+
non-contiguous input shifts the output by ~1.0), so `forward` forces a
64+
contiguous tensor. With these two fixes the encoder reproduces **bit-exactly**
65+
across processes on the same env + GPU.
66+
67+
## Input / output
68+
69+
| tensor | axes | notes |
70+
| --- | --- | --- |
71+
| `input` | `b, c(3: red/green/blue), y=1008, x=1008` | raw RGB in `[0, 255]`; scale-to-[0,1] + ImageNet mean/std happen inside the wrapper |
72+
| `vision_features` | `b, c(256), y=72, x=72` | SAM 3 backbone top-level FPN feature map |
73+
74+
## Reproducing this export
75+
76+
With `sam3` installed and the checkpoint available (see [Setup](#setup)):
77+
78+
```console
79+
python sam3_original.py # run the original encoder, save reference tensors
80+
python bioimageio_export.py # build, export and test the bioimage.io package
81+
```
82+
83+
- `sam3_encoder_bioimageio.py` — the bundled `Sam3EncoderBioImageIO` wrapper the
84+
exported model runs against (stock, unmodified `sam3`).
85+
- `bioimageio_export.py` writes the packaged model and `output/environment.yaml`.

0 commit comments

Comments
 (0)