From d0a88feff99bb333f4fe9b15dc7369c425fe85d4 Mon Sep 17 00:00:00 2001 From: belkassaby Date: Fri, 22 May 2026 15:08:27 -0400 Subject: [PATCH] update documentation --- README.md | 2 +- docs/HF-MODEL-CARD.md | 15 ++-- docs/PLAN.md | 159 +----------------------------------------- 3 files changed, 13 insertions(+), 163 deletions(-) diff --git a/README.md b/README.md index 8e99f22..5a91e5c 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/docs/HF-MODEL-CARD.md b/docs/HF-MODEL-CARD.md index e0cfb8f..f620d2e 100644 --- a/docs/HF-MODEL-CARD.md +++ b/docs/HF-MODEL-CARD.md @@ -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 @@ -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 diff --git a/docs/PLAN.md b/docs/PLAN.md index c5671a2..0174136 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -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 @@ -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 @@ -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. ---