Skip to content

update documentation #7

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
belkassaby merged 1 commit into from
May 22, 2026
Merged

update documentation #7

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@

Browser-side cellular segmentation powered by [Cellpose-SAM](https://github.com/MouseLand/cellpose), running on WebGPU. Faithful TypeScript port of the Cellpose-SAM inference + dynamics pipeline, designed for in-browser microscopy workflows without a server round-trip.

> **Status:** v0.1.0 — first end-to-end-working release. Phase 1 from the [implementation plan](./docs/PLAN.md) is complete: model loading + IndexedDB cache, preprocessing, WebGPU inference in a worker, tile averaging, flow dynamics, full-image label maps. Phase 2 (SlimSAM-style compression + domain-specialized slim models) is planned but not started.
> **Status:** v0.1.0 — first end-to-end-working release. The full port from the [implementation plan](./docs/PLAN.md) is complete: model loading + IndexedDB cache, preprocessing, WebGPU inference in a worker, tile averaging, flow dynamics, full-image label maps. SlimSAM-style compression and domain-specialized finetunes are out of scope — see the plan's §6 for the rationale.

## Highlights

Expand Down
15 changes: 9 additions & 6 deletions docs/HF-MODEL-CARD.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,8 +48,9 @@ This is the **same model** as `mouseland/cellpose-sam`, just re-packaged as a
no pruning, no architecture changes. Numerical parity vs the PyTorch source is
**1.24e-05 worst max-abs-error** across 10 random tiles (gate 1e-3).

If you want a smaller model, this is not it — that's Phase 2 (SlimCPSAM),
which has not been built yet.
If you want a smaller model, this package isn't it — a SlimSAM-style
compressed CPSAM would be a multi-month research effort and is not on the
roadmap.

## Intended use

Expand Down Expand Up @@ -223,10 +224,12 @@ dynamics.
is non-negotiable for FP16 graph IO under ORT-web 1.20+.
- **Not promptable.** This is dense per-pixel regression, not SAM-style
prompted segmentation.
- **No domain-specialized variants.** This is the generalist CPSAM. Historical
Cellpose variants (cyto / cyto2 / cyto3 / nuclei) and Omnipose bacteria
models are not provided here — they are planned as **Phase 2** SlimCPSAM
finetunes and have not been built yet.
- **No separately-packaged domain variants.** This is the generalist CPSAM.
Per Pachitariu et al. 2025, the base CPSAM weights are already trained on
the major public segmentation corpora (Cellpose 1/2/3 cyto+nuclei, LIVECell,
TissueNet, DeepBacs, Omnipose bact_phase/bact_fluor, NeurIPS 2022
multimodality challenge), so re-packaging a per-domain finetune from those
same datasets would be redundant.
- **No INT8 / FP32 fallback.** WebGPU only.

## Citation
Expand Down
159 changes: 3 additions & 156 deletions docs/PLAN.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,10 @@ Browser-side Cellpose-SAM via a new `Cellpose.js` package, consumed by jit-ui as

**Scope decisions (locked in):**

- **Phase 1 only ships first** — port stock CPSAM to the browser. No SlimSAM-style compression in Phase 1.
- **Scope is port-only** — port stock CPSAM to the browser. No SlimSAM-style compression, no domain-specialized finetunes.
- **WebGPU required** — ship FP16, no WASM fallback in v1.
- **Separate repo** `Cellpose.js` — not an Nx lib inside jit-ui.

**Phase 2 planned but NOT started** — SlimSAM-style compression of CPSAM (§7) plus domain-specialized slim finetunes (cyto/cyto2/cyto3, nuclei, Omnipose bacteria, TissueNet, LiveCell). Research/ML project, multi-month, gated on Phase 1 shipping. Do not begin without explicit go.

---

## Phase 1 status (as of 2026-05-15): COMPLETE — v0.1.1
Expand Down Expand Up @@ -99,7 +97,7 @@ From `cellpose.vit_sam` and the Cellpose-SAM preprint (bioRxiv 2025.04.28):

- SlimSAM (Chen et al., NeurIPS 2024) is **channel-pruned + distilled SAM-ViT-B** plus the _standard_ SAM prompt encoder and SAM mask decoder. `Xenova/slimsam-50-uniform` is the ONNX export already wired into Transformers.js (jit-ui already uses the sibling `Xenova/sam-vit-base` in `transformers.worker.ts:238`).
- It is **promptable mask generation**, not dense regression. Outputs are `pred_masks` + `iou_scores`, given `input_points` / `input_boxes`.
- It is **not a drop-in replacement** for CPSAM's head. Plugging SlimSAM in where CPSAM goes would produce per-prompt mask candidates, not flow-derived instance masks. Making a SlimSAM-style CPSAM would require either (a) pruning CPSAM's ViT-L while keeping its flow head, or (b) training a new flow head on top of SlimSAM's encoder. Both are multi-month ML work, out of scope for Phase 1.
- It is **not a drop-in replacement** for CPSAM's head. Plugging SlimSAM in where CPSAM goes would produce per-prompt mask candidates, not flow-derived instance masks. Making a SlimSAM-style CPSAM would require either (a) pruning CPSAM's ViT-L while keeping its flow head, or (b) training a new flow head on top of SlimSAM's encoder. Both are multi-month ML work, out of scope for this package.

### 1.3 Specific feasibility checks

Expand Down Expand Up @@ -337,158 +335,7 @@ Touchpoints in this repo:
- INT8 path. Revisit only if FP16 download size becomes a complaint.
- WASM fallback. WebGPU required.
- **Pre-Chrome-135 / pre-Safari-17.4 browsers.** Native `Float16Array` is mandatory for FP16 graph IO. Reaching older browsers would require a working post-export FP16 conversion path (currently broken, see §1.5) or an FP32 model (~1.1 GB, double the bandwidth).
- SlimSAM-style compression of CPSAM. Separate research project; not part of Phase 1.

---

## 7. Phase 2 — SlimCPSAM and Domain-Specialized Slim Models (PLANNED, NOT STARTED)

A research workstream to produce small (~50–150 MB) browser-runnable cellular
segmentation models. Gated on Phase 1 shipping. Multi-month timeline. Requires
GPU compute, labeled data access, and ML iteration. **Do not begin without an
explicit go decision and a compute/budget allocation.**

### 7.1 Reframing the model list

The list of targets needs cleanup before planning. The original ask combined
architectures, model variants, and datasets:

| Item | What it actually is | SlimSAM strategy |
| ---------------------- | --------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------- |
| **CPSAM** | ViT-L + flow head (Cellpose 4 default) | **Compress directly** — channel-prune + distill ViT-L |
| **cyto, cyto2, cyto3** | Pre-v4 **U-Net** generalist cytoplasm models (Cellpose 1/2/3) | Re-train slim CPSAM head on cyto-style data — _not_ prune U-Net |
| **nuclei** | Pre-v4 U-Net nuclei model | Same as cyto: finetune slim CPSAM on nuclei corpus |
| **Bacteria (omni)** | Omnipose `bact_phase_omni` / `bact_fluor_omni` — U-Net + Omnipose distance-transform postprocessing | Finetune slim CPSAM on bacteria images; keep CPSAM flow postprocessing (no need to port Omnipose dynamics) |
| **Bacteria (omni+)** | Omnipose's improved bacteria models (`bact_phase_cp`, `bact_phase_omnipose+`) | Same as above |
| **TissueNet** | **Dataset**, not a model. Multi-channel tissue imaging. | Finetune slim CPSAM on TissueNet |
| **LiveCell** | **Dataset**, not a model. Phase-contrast cell line images. | Finetune slim CPSAM on LiveCell |

Key insight: only CPSAM itself is a SAM-derived architecture. Everything else
in the list is either a different architecture (U-Net) or a dataset. So the
strategy is **not** "make a SlimSAM version of each existing model" — it is
"build one SlimCPSAM (the architecture-compression step), then produce 7
domain-specialized finetunes from that single slim backbone." This gives us
the same domain coverage as historical Cellpose without porting U-Net code or
Omnipose's distance-transform postprocessing.

### 7.2 Compute, data, and timeline preconditions

| Resource | Need | Notes |
| ------------------- | ---------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| GPU compute | ~500–1500 GPU-hours | A100/H100 preferred; a 4090 desktop is the floor (slower). The pruning-distill loop dominates. |
| Labeled data | Public training corpora for each domain | TissueNet (gated — research registration), LiveCell (public), Omnipose bacteria (public via Cutler lab), Cellpose 4 training set (public via the Cellpose-SAM release), cyto/nuclei legacy (public via Cellpose 1–3 release) |
| Storage | ~500 GB | Datasets + checkpoints + intermediate distillation states |
| Eval infrastructure | Per-domain held-out test splits + matching metrics (per-instance IoU, AP@0.5, segmentation accuracy) | Reuse Cellpose's evaluation harness where possible |
| Personnel | One ML engineer focused for ~3 months, or shared across ~5 months | The pruning recipe is non-trivial; expect iteration |
| Cost | $5–20K in cloud GPU if not done on owned hardware | Lower if 4090 desktop is used end-to-end |

If any of these is missing, **the right call is to skip Phase 2** and rely on
Phase 1's full-precision CPSAM in the browser.

### 7.3 Sub-phase A — Build SlimCPSAM (the base slim model)

The one-time architecture-compression step. Output: a single ~50–150 MB ONNX
model with the same input/output contract as CPSAM but a pruned ViT backbone.

**Approach.** Apply SlimSAM's recipe (Chen et al., NeurIPS 2024) to CPSAM's
ViT-L encoder, keeping CPSAM's flow regression head. SlimSAM was developed
against SAM-ViT-B, so the recipe needs adaptation for ViT-L (longer training,
adjusted pruning ratios). The flow head is small (~2 MB) and stays untouched
structurally but may need finetuning after backbone pruning.

**Steps**

1. **Replicate SlimSAM training infra.** Reproduce the public SlimSAM training
pipeline on a small SAM-ViT-B sanity check before touching CPSAM. Goal:
confirm we can hit the published SlimSAM-50 IoU numbers within 2 points.
Time-box: 2 weeks.
2. **Adapt to ViT-L + CPSAM head.** Modify the SlimSAM training loop:
- Teacher: full CPSAM (ViT-L, 304 M params).
- Student: pruned ViT-L (target ~50% channels retained on ViT-S/B-class
scale, OR retrain a ViT-B-sized student from scratch with SlimSAM-style
distillation).
- Loss: SlimSAM's intermediate feature alignment + a regression loss on the
CPSAM flow outputs (`flow_y`, `flow_x`, `cellprob`).
- Training data: Cellpose 4's public training corpus (the same data CPSAM
was trained on).
Time-box: 4 weeks of training iteration.
3. **Quality gate.** SlimCPSAM mean per-instance IoU on the Cellpose 4
validation set must be **within 3 points** of full CPSAM. If we can't hit
that, abandon Phase 2 — the slim model is not worth a domain-specific
training campaign that bakes in the quality gap.
4. **Export to ONNX FP16**, single-file, same IO contract as Phase 1's CPSAM
ONNX so `cellpose-js` can swap models transparently.

**Deliverable**: `ballon999/slimcpsam-onnx` on HF Hub (matches the Phase 1 account `ballon999/cellpose-sam-onnx`), ~50–150 MB.

### 7.4 Sub-phase B — Domain-specialized slim finetunes

Once SlimCPSAM exists and clears the quality gate, finetune it on each domain
corpus.

| Slim model | Training corpus | Source | License notes |
| ---------------------- | ------------------------------------------------------------ | ---------------------------- | ------------------------------------------- |
| `slimcpsam-cyto` | Cellpose 1 cyto training data | MouseLand `cellpose` release | BSD-3 |
| `slimcpsam-cyto2` | Cellpose 2 cyto2 training data | MouseLand `cellpose` release | BSD-3 |
| `slimcpsam-cyto3` | Cellpose 3 cyto3 training data | MouseLand `cellpose` release | BSD-3 |
| `slimcpsam-nuclei` | Cellpose nuclei training data | MouseLand `cellpose` release | BSD-3 |
| `slimcpsam-bact-omni` | Omnipose `bact_phase_omni` + `bact_fluor_omni` training data | kcutler `omnipose` release | MIT, with Cutler-lab attribution |
| `slimcpsam-bact-omni+` | Omnipose's improved bacteria corpus | kcutler `omnipose` release | MIT |
| `slimcpsam-tissuenet` | TissueNet | vanvalenlab (Caltech) | Research-use license; registration required |
| `slimcpsam-livecell` | LiveCell | Sartorius / Chalmers | CC-BY-NC 4.0 (non-commercial!) |

**Strategy per finetune**

1. **Frozen backbone + head finetune (1–3 epochs).** Cheap, fast, often
sufficient if the domain isn't far from CPSAM's training distribution.
2. **Full finetune (5–10 epochs)** only if the frozen-backbone pass doesn't
meet a per-domain quality gate (e.g., AP@0.5 within 2 points of the
original full-precision Cellpose/Omnipose model on that domain's test set).
3. **Per-domain evaluation against two baselines:**
- The historical Cellpose/Omnipose model for that domain (the user-facing
comparison: "is the slim version still good enough?").
- Full CPSAM (the upper-bound comparison: "how much accuracy are we trading
for size?").
4. **Re-export to ONNX FP16** with the same IO contract; publish to HF Hub.

**Effort per model**: ~3–5 days of training + eval, parallelizable across GPUs.
Total: ~6–8 weeks if run serially, ~2–3 weeks with 3+ GPUs in parallel.

**License gotcha — LiveCell.** LiveCell's CC-BY-NC 4.0 license restricts
commercial use. If jit-ui or downstream Jackson Lab products have any
commercial use, `slimcpsam-livecell` must be flagged as non-commercial-only
in the package metadata, or the dataset must be replaced. Decide before
training, not after.

### 7.5 Shipping plan in `cellpose-js` and jit-ui

After Phase 2, the `cellpose-js` API extends to support model selection:

```ts
const cp = await Cellpose.fromPretrained('slimcpsam-cyto3', { preload: true });
// or
const cp = await Cellpose.fromPretrained('cpsam', { preload: true }); // Phase 1 full-precision
```

In jit-ui's pipeline-dialog operation descriptor, add a `model` dropdown to the
`cellpose-segment` op with the 8 slim variants plus the full CPSAM. Default:
`slimcpsam-cyto3` (smallest acceptable generalist), with full CPSAM available
for users who want maximum accuracy and can wait for the 588 MB download.

Caching footprint: if a user runs all 9 models, ~9 × ~100 MB = ~900 MB total
in IndexedDB. Add a "clear cached models" affordance in jit-ui settings.

### 7.6 Decision points (revisit before starting Phase 2)

1. **Does Phase 1 generalize well enough?** If users find full CPSAM works for
every domain they care about, the per-domain finetunes are wasted effort.
Collect user feedback after Phase 1 ships before committing to §7.4.
2. **Compute availability.** Without ≥1 A100-class GPU for ~6 weeks, Phase 2
stalls. Confirm before kickoff.
3. **TissueNet / LiveCell access and licensing.** Datasets gated or
non-commercial — confirm acceptable use up front, not mid-training.
4. **SlimCPSAM quality gate (§7.3 step 3).** If we can't hit "within 3 IoU
points of full CPSAM," abandon. Don't ship a strictly worse smaller model.
- SlimSAM-style compression of CPSAM. Not part of v1; no follow-up project planned. CPSAM is already trained on the major public corpora (Cellpose 1/2/3, LIVECell, TissueNet, DeepBacs, Omnipose, NeurIPS 2022) per Pachitariu et al. 2025, so domain-specialized finetunes on those datasets are no-ops.

---

Expand Down
Loading