diff --git a/CHANGELOG.md b/CHANGELOG.md index aa1fd01..b8074ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,25 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.3.0] — 2026-06-14 + +### Changed + +- **Full pipeline runs in the Web Worker.** `segment()` now offloads the + entire pipeline — preprocessing, per-tile inference, tile averaging, flow + dynamics, and inverse-resize — to the inference worker via a single + `segment` message, and transfers back only the final label map. Previously + only per-tile inference ran off-thread, so preprocessing and dynamics ran on + the main thread and blocked the UI. The public `segment(input, opts)` API and + `onTileProgress` are unchanged. + +### Removed + +- `SegmentOutput.tiles[].flows_cellprob` is now an empty `Float32Array` — the + per-tile flow tensors stay in the worker (returning them would defeat the + offload). The per-tile timing diagnostics (`tx`/`ty`/`bsize`/`inferenceMs`) + are unchanged. + ## [0.2.0] — 2026-05-22 Python-parity fixes plus quality-of-life additions surfaced by a code review diff --git a/README.md b/README.md index 89600c5..a51e466 100644 --- a/README.md +++ b/README.md @@ -8,14 +8,14 @@ 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. 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. +> **Status:** The full pipeline — preprocessing, WebGPU inference, tile averaging, and flow dynamics — runs in a Web Worker, so `segment()` never blocks the UI thread (earlier releases ran only per-tile inference off-thread). Model loading uses an IndexedDB cache. SlimSAM-style compression and domain-specialized finetunes are out of scope — see the [implementation plan](./docs/PLAN.md) §6. ## Highlights - **Single-call API**: `await Cellpose.fromPretrained(modelUrl)` → `await cp.segment(image, opts)` → a `Uint32Array` instance label map at source resolution. - **WebGPU inference** via `onnxruntime-web/webgpu`. Measured **~277 ms / 256×256 tile on an M1 Max**. Cold start ~2.3 s (one-time shader compile). -- **Web Worker offload**: inference doesn't block the UI thread; AbortSignal terminates the worker mid-run with sub-100 ms latency. -- **Faithful Python parity** for preprocess and dynamics — 14/14 vitest parity tests pass against numpy-generated `.npy` fixtures. +- **Fully off the UI thread**: the entire pipeline (preprocess → inference → tile averaging → flow dynamics → resize) runs in a Web Worker via a single `segment` message; only the final label map is transferred back, so the UI never blocks. `AbortSignal` terminates the worker mid-run with sub-100 ms latency. +- **Faithful Python parity** for preprocess and dynamics — vitest parity tests pass against numpy-generated `.npy` fixtures. - **IndexedDB cache** for the 588 MB FP16 model: first visit fetches from your CDN; subsequent visits load from local storage in <2 s. ## Browser requirements @@ -72,7 +72,8 @@ console.log(`Found ${result.count} cells.`); // result.width : number — source image width // result.height : number — source image height // result.totalMs : number — wall-clock time for the segment() call -// result.tiles : per-tile diagnostics (flow tensors, inference time) +// result.tiles : per-tile timing diagnostics (tx/ty/bsize/inferenceMs). +// Flow tensors stay in the worker, so flows_cellprob is empty. ``` ## Parameter quick-reference @@ -113,9 +114,11 @@ Rescales the image so the median cell occupies ~30 px (CPSAM's training median). ## Architecture +The whole pipeline runs **inside the Web Worker**: the main thread posts the image with one `segment` message and receives the final `masks` (transferred back), so preprocessing and flow dynamics never block the UI. + ``` input image → buildCpsamChannels → diameterResize → normalizePerChannel → makeTiles - ⇣ (per tile, via worker) + ⇣ (per tile) ort.InferenceSession.run ⇣ averageTiles @@ -132,7 +135,7 @@ See [`src/`](./src/) for module-level documentation. ## Testing ```sh -npm run test # vitest: 14 parity tests against numpy fixtures +npm run test # vitest: parity + unit tests against numpy fixtures npm run typecheck # tsc --noEmit npm run build # vite library build + tsc --emitDeclarationOnly npm run demo # vite serve examples/demo diff --git a/package.json b/package.json index e3b7ad8..015d27d 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "cellpose-js", - "version": "0.2.0", + "version": "0.3.0", "description": "Browser-side cellular segmentation via Cellpose-SAM, running on WebGPU.", "type": "module", "license": "MIT", diff --git a/src/cellpose.ts b/src/cellpose.ts index 4a1d1ac..32c200d 100644 --- a/src/cellpose.ts +++ b/src/cellpose.ts @@ -1,21 +1,16 @@ /** - * Public Cellpose API. Milestone 3: inference offloaded to a Web Worker, - * AbortSignal-driven cancellation, tile-level progress. + * Public Cellpose API. The full segmentation pipeline (preprocess → tile → + * inference → average → dynamics → inverse-resize) runs in a Web Worker via the + * `segment` message, so nothing heavy runs on the UI thread. AbortSignal-driven + * cancellation and tile-level progress are supported. The legacy per-tile + * `_runTile` path is retained for back-compat. */ import { assertSupportedEnvironment, describeAdapter } from './env.js'; import { fetchModel, type FetchProgress } from './model-cache.js'; import { configureOrt, _getWasmPaths } from './session.js'; -import { - buildCpsamChannels, - type ChannelMapOptions, - diameterResize, - normalizePerChannel, - type NormalizeOptions, - makeTiles, - type TileRecord, -} from './preprocess/index.js'; -import { computeMasks, type ComputeMasksOptions, averageTiles } from './postprocess/index.js'; -import type { MainToWorker, WorkerToMain } from './worker-protocol.js'; +import { type ChannelMapOptions, type NormalizeOptions } from './preprocess/index.js'; +import { type ComputeMasksOptions } from './postprocess/index.js'; +import type { MainToWorker, WorkerToMain, WorkerSegmentOptions } from './worker-protocol.js'; export interface FromPretrainedOptions { /** Eagerly create the inference worker + ORT session at construct time. */ @@ -71,8 +66,8 @@ export interface SegmentOutput { /** Width / height at source resolution. */ width: number; height: number; - /** Per-tile diagnostics. Per-tile masks are NOT included (single global label - * map is the v1 contract). */ + /** Per-tile diagnostics (timing). Flow tensors stay in the worker, so + * `flows_cellprob` is empty here — returning them would defeat the offload. */ tiles: SegmentTileOutput[]; /** Resized image dimensions (intermediate; pre-inverse-resize). */ resizedWidth: number; @@ -98,12 +93,20 @@ interface PendingTile { reject: (err: Error) => void; } +interface PendingSegment { + resolve: (out: SegmentOutput) => void; + reject: (err: Error) => void; + onTileProgress?: (done: number, total: number) => void; + t0: number; +} + export class Cellpose { private _worker: Worker | null = null; private _workerReady: Promise | null = null; private _adapterInfo: { vendor: string; architecture: string; device: string } | null = null; - private _nextTileId = 0; private _pending = new Map(); + private _nextReqId = 0; + private _pendingSegments = new Map(); // _modelBytes is detached after the worker takes ownership of it. To respawn // after _abort() we re-fetch via fetchModel (cache hit -> instant). @@ -162,8 +165,8 @@ export class Cellpose { }); this._worker = worker; // Track whether the worker has finished init. Before this flips true, - // a worker `error` event has nowhere to be reported (pending-tile queue - // is empty), so we route it to the init promise's reject instead. The + // a worker `error` event has nowhere to be reported (pending queues are + // empty), so we route it to the init promise's reject instead. The // ref is updated to a no-op once init completes. let rejectInit: (err: Error) => void = () => {}; let initSettled = false; @@ -186,6 +189,35 @@ export class Cellpose { } else if (!initSettled) { rejectInit(new Error(msg.message)); } + } else if (msg.type === 'segment-progress') { + // Only per-tile inference maps onto the public onTileProgress callback; + // the dynamics phase is signalled by the final (done===total) tile tick. + if (msg.phase === 'inference') { + this._pendingSegments.get(msg.reqId)?.onTileProgress?.(msg.done, msg.total); + } + } else if (msg.type === 'segment-result') { + const p = this._pendingSegments.get(msg.reqId); + if (p) { + this._pendingSegments.delete(msg.reqId); + p.resolve({ + masks: msg.masks, + count: msg.count, + width: msg.width, + height: msg.height, + tiles: msg.tiles.map((t) => ({ ...t, flows_cellprob: new Float32Array(0) })), + resizedWidth: msg.resizedWidth, + resizedHeight: msg.resizedHeight, + scale: msg.scale, + totalMs: performance.now() - p.t0, + postprocessMs: msg.postprocessMs, + }); + } + } else if (msg.type === 'segment-error') { + const p = this._pendingSegments.get(msg.reqId); + if (p) { + this._pendingSegments.delete(msg.reqId); + p.reject(new Error(msg.message)); + } } else if (msg.type === 'status') { opts.onStatus?.(msg.status); } @@ -199,8 +231,7 @@ export class Cellpose { if (!initSettled) { rejectInit(new Error(`worker module failed to load: ${detail}`)); } - for (const p of this._pending.values()) p.reject(new Error(detail)); - this._pending.clear(); + this._rejectAllPending(new Error(detail)); }); // messageerror fires when structured-clone deserialization fails on a // message *received* by the worker (the failure is reported on the main @@ -208,8 +239,7 @@ export class Cellpose { worker.addEventListener('messageerror', () => { const err = new Error('worker received malformed message (structured-clone failure)'); if (!initSettled) rejectInit(err); - for (const p of this._pending.values()) p.reject(err); - this._pending.clear(); + this._rejectAllPending(err); }); this._workerReady = new Promise((resolve, reject) => { @@ -278,35 +308,25 @@ export class Cellpose { return this._workerReady; } - /** Aborts the worker mid-run; pending tile promises reject with AbortError. */ + /** Reject every in-flight tile + segment promise (worker died/errored). */ + private _rejectAllPending(err: Error): void { + for (const p of this._pending.values()) p.reject(err); + this._pending.clear(); + for (const p of this._pendingSegments.values()) p.reject(err); + this._pendingSegments.clear(); + } + + /** Aborts the worker mid-run; pending promises reject with AbortError. */ private _abort(reason?: string): void { if (!this._worker) return; this._worker.terminate(); this._worker = null; this._workerReady = null; - const err = new DOMException(reason ?? 'Operation aborted', 'AbortError'); - for (const p of this._pending.values()) p.reject(err); - this._pending.clear(); - } - - private async _runTile( - tile: Float32Array, - bsize: number, - ): Promise> { - await this._ensureWorker(); - const worker = this._worker; - if (!worker) throw new Error('worker not available'); - const tileId = this._nextTileId++; - return new Promise>((resolve, reject) => { - this._pending.set(tileId, { resolve, reject }); - const msg: MainToWorker = { type: 'run-tile', tileId, tile, bsize }; - worker.postMessage(msg, [tile.buffer]); - }); + this._rejectAllPending(new DOMException(reason ?? 'Operation aborted', 'AbortError')); } async segment(input: SegmentInput, opts: SegmentOptions = {}): Promise { const t0 = performance.now(); - const tileSize = opts.tile ?? 256; const signal = opts.signal; // Hook abort: terminate the worker, surface AbortError to the caller. @@ -319,96 +339,36 @@ export class Cellpose { } try { - let chw = buildCpsamChannels(input.data, input.width, input.height, input.channels, opts); - let w = input.width, - h = input.height, - scale = 1; - // Match cellpose's order: normalize FIRST, then resize. (See - // models.py:_run_net: `transforms.normalize_img(x)` then - // `transforms.resize_image(imgi[b], rsz=rsz)`.) Resizing first would - // change the per-channel percentile statistics that drive normalization. - chw = normalizePerChannel(chw, 3, w * h, opts.normalize ?? {}); - if (opts.diameter !== undefined) { - const r = diameterResize(chw, w, h, { channels: 3, diameter: opts.diameter }); - chw = r.data; - w = r.width; - h = r.height; - scale = r.scale; - } - - const tileOpts: { bsize: number; overlap?: number } = { bsize: tileSize }; - if (opts.overlap !== undefined) tileOpts.overlap = opts.overlap; - const tiles: TileRecord[] = makeTiles(chw, w, h, 3, tileOpts); - - const out: SegmentTileOutput[] = []; - for (let i = 0; i < tiles.length; i++) { - const t = tiles[i]!; - const r = await this._runTile(t.tile, tileSize); - out.push({ - flows_cellprob: r.flowsCellprob, - tx: t.tx, - ty: t.ty, - bsize: tileSize, - inferenceMs: r.inferenceMs, - }); - opts.onTileProgress?.(i + 1, tiles.length); - } + await this._ensureWorker(); + const worker = this._worker; + if (!worker) throw new Error('worker not available'); - // Average overlapping tile predictions into a single full-image tensor, - // then run dynamics once. See average_tiles.ts header for why this is the - // right algorithm choice (vs per-tile dynamics + label stitching). - const tPost = performance.now(); - const averaged = averageTiles( - out.map((o) => ({ flowsCellprob: o.flows_cellprob, tx: o.tx, ty: o.ty, bsize: o.bsize })), - h, - w, - ); - const hwFull = h * w; - const dPFull = averaged.data.subarray(0, 2 * hwFull) as Float32Array; - const cpFull = averaged.data.subarray(2 * hwFull, 3 * hwFull) as Float32Array; - // Match cellpose's niter scaling: `niter = int(200 / image_scaling)` - // where image_scaling = 30 / diameter (== JS's `scale`). Upscaled images - // need fewer iterations (each step covers more source pixels) and - // downscaled images need more. Honors any explicit user override. - const dynOpts: ComputeMasksOptions = { ...(opts.dynamics ?? {}) }; - if (dynOpts.niter === undefined && opts.diameter !== undefined && scale !== 1) { - dynOpts.niter = Math.max(1, Math.floor(200 / scale)); - } - const m = computeMasks(dPFull, cpFull, h, w, dynOpts); + // Strip the non-serializable opts (AbortSignal + callbacks); the rest is + // structured-cloned to the worker. + const { signal: _omitSignal, onTileProgress, ...rest } = opts; + void _omitSignal; + const workerOpts = rest as WorkerSegmentOptions; + const reqId = this._nextReqId++; - // Inverse-resize labels back to source resolution if a diameter resize - // was applied (nearest-neighbor — labels must not be interpolated). - let masksSrc = m.masks; - let outW = w, - outH = h; - if (scale !== 1) { - outW = input.width; - outH = input.height; - const resized = new Uint32Array(outW * outH); - for (let yy = 0; yy < outH; yy++) { - const sy = Math.min(h - 1, Math.max(0, Math.round(yy * scale))); - const srcRow = sy * w; - const dstRow = yy * outW; - for (let xx = 0; xx < outW; xx++) { - const sx = Math.min(w - 1, Math.max(0, Math.round(xx * scale))); - resized[dstRow + xx] = m.masks[srcRow + sx] as number; - } - } - masksSrc = resized; - } - const postprocessMs = performance.now() - tPost; - return { - masks: masksSrc, - count: m.count, - width: outW, - height: outH, - tiles: out, - resizedWidth: w, - resizedHeight: h, - scale, - totalMs: performance.now() - t0, - postprocessMs, - }; + return await new Promise((resolve, reject) => { + const pending: PendingSegment = { resolve, reject, t0 }; + if (onTileProgress) pending.onTileProgress = onTileProgress; + this._pendingSegments.set(reqId, pending); + const msg: MainToWorker = { + type: 'segment', + reqId, + image: { + data: input.data, + width: input.width, + height: input.height, + channels: input.channels, + }, + opts: workerOpts, + }; + // image.data is cloned (left intact for the caller); the result masks + // are transferred back from the worker. + worker.postMessage(msg); + }); } finally { if (signal && abortListener) signal.removeEventListener('abort', abortListener); } @@ -421,5 +381,6 @@ export class Cellpose { this._worker = null; this._workerReady = null; this._pending.clear(); + this._pendingSegments.clear(); } } diff --git a/src/inference.worker.ts b/src/inference.worker.ts index 4061ea0..c41b1ff 100644 --- a/src/inference.worker.ts +++ b/src/inference.worker.ts @@ -1,19 +1,38 @@ /** - * Inference worker — hosts the ORT-WebGPU session so per-tile inference - * (~280 ms each) doesn't block the UI thread. + * Inference worker — hosts the ORT-WebGPU session AND (via the `segment` + * message) runs the entire Cellpose pipeline off the UI thread: preprocess → + * tile → per-tile inference → tile averaging → flow dynamics → inverse-resize. + * Only the final label map is posted back. This keeps the main thread free — + * preprocess and dynamics used to run there and froze the UI. * * Lifecycle: - * main thread: new Worker(...) -> postMessage({type:'init', modelBytes}) - * worker: ORT session create -> postMessage({type:'ready', adapterInfo}) - * main thread: postMessage({type:'run-tile', tileId, tile, bsize}, [tile.buffer]) - * worker: sess.run(...) -> postMessage({type:'tile-result', tileId, flowsCellprob}, [flowsCellprob.buffer]) + * main: new Worker(...) -> postMessage({type:'init', modelBytes}) + * worker: ORT session create -> postMessage({type:'ready', adapterInfo}) + * main: postMessage({type:'segment', reqId, image, opts}) + * worker: postMessage({type:'segment-progress', ...})* -> + * postMessage({type:'segment-result', reqId, masks, ...}, [masks.buffer]) + * + * The legacy `run-tile` message (single-tile inference, main thread orchestrates) + * is kept for back-compat. * * Cancellation: main thread calls worker.terminate(); worker dies mid-run. - * No graceful unwind needed — the session is gone with the worker. */ /// import * as ort from 'onnxruntime-web/webgpu'; import type { MainToWorker, WorkerToMain } from './worker-protocol.js'; +import { + buildCpsamChannels, + normalizePerChannel, + diameterResize, + makeTiles, + type TileRecord, +} from './preprocess/index.js'; +import { + averageTiles, + computeMasks, + type ComputeMasksOptions, + type TileInputForAveraging, +} from './postprocess/index.js'; let session: ort.InferenceSession | null = null; let wasmPathsConfigured = false; @@ -59,33 +78,24 @@ async function handleInit(msg: Extract): Promise postReply({ type: 'ready', adapterInfo }); } -async function handleRunTile(msg: Extract): Promise { - const sess = session; - if (!sess) { - postReply({ type: 'error', tileId: msg.tileId, message: 'worker not initialized' }); - return; - } - +/** Run one (C,B,B) tile through the model; returns FP32 flows+cellprob. */ +async function runTile( + sess: ort.InferenceSession, + tile: Float32Array, + bsize: number, +): Promise<{ flowsCellprob: Float32Array; inferenceMs: number }> { // FP32 -> FP16. Float16Array auto-rounds on store. - const fp16 = new Float16Array(msg.tile.length); - for (let i = 0; i < msg.tile.length; i++) fp16[i] = msg.tile[i] as number; - // ort-web 1.26's TS types don't yet accept Float16Array directly; cast. - const tensor = new ort.Tensor('float16', fp16 as unknown as Uint16Array, [ - 1, - 3, - msg.bsize, - msg.bsize, - ]); + const fp16 = new Float16Array(tile.length); + for (let i = 0; i < tile.length; i++) fp16[i] = tile[i] as number; + // ort-web's TS types don't yet accept Float16Array directly; cast. + const tensor = new ort.Tensor('float16', fp16 as unknown as Uint16Array, [1, 3, bsize, bsize]); const t0 = performance.now(); const outputs = await sess.run({ [sess.inputNames[0] as string]: tensor }); const inferenceMs = performance.now() - t0; const out = outputs['flows_cellprob']; - if (!out) { - postReply({ type: 'error', tileId: msg.tileId, message: `missing 'flows_cellprob' output` }); - return; - } + if (!out) throw new Error(`missing 'flows_cellprob' output`); // FP16 -> FP32. Reinterpret the Uint16Array as Float16Array on the same // buffer, then copy into FP32. @@ -93,10 +103,120 @@ async function handleRunTile(msg: Extract): const outF16 = new Float16Array(u16.buffer, u16.byteOffset, u16.length); const outF32 = new Float32Array(outF16.length); for (let i = 0; i < outF16.length; i++) outF32[i] = outF16[i] as number; + return { flowsCellprob: outF32, inferenceMs }; +} - postReply({ type: 'tile-result', tileId: msg.tileId, flowsCellprob: outF32, inferenceMs }, [ - outF32.buffer, - ]); +async function handleRunTile(msg: Extract): Promise { + const sess = session; + if (!sess) { + postReply({ type: 'error', tileId: msg.tileId, message: 'worker not initialized' }); + return; + } + const r = await runTile(sess, msg.tile, msg.bsize); + postReply( + { + type: 'tile-result', + tileId: msg.tileId, + flowsCellprob: r.flowsCellprob, + inferenceMs: r.inferenceMs, + }, + [r.flowsCellprob.buffer], + ); +} + +/** Full pipeline in the worker — keeps preprocess + dynamics off the UI thread. */ +async function handleSegment(msg: Extract): Promise { + const sess = session; + if (!sess) { + postReply({ type: 'segment-error', reqId: msg.reqId, message: 'worker not initialized' }); + return; + } + const { reqId, image, opts } = msg; + const tileSize = opts.tile ?? 256; + + // Preprocess: channels -> normalize -> (optional) diameter resize. + let chw = buildCpsamChannels(image.data, image.width, image.height, image.channels, opts); + let w = image.width, + h = image.height, + scale = 1; + chw = normalizePerChannel(chw, 3, w * h, opts.normalize ?? {}); + if (opts.diameter !== undefined) { + const r = diameterResize(chw, w, h, { channels: 3, diameter: opts.diameter }); + chw = r.data; + w = r.width; + h = r.height; + scale = r.scale; + } + + const tileOpts: { bsize: number; overlap?: number } = { bsize: tileSize }; + if (opts.overlap !== undefined) tileOpts.overlap = opts.overlap; + const tiles: TileRecord[] = makeTiles(chw, w, h, 3, tileOpts); + + const flows: TileInputForAveraging[] = []; + const tileDiags: { tx: number; ty: number; bsize: number; inferenceMs: number }[] = []; + for (let i = 0; i < tiles.length; i++) { + const t = tiles[i]!; + const r = await runTile(sess, t.tile, tileSize); + flows.push({ flowsCellprob: r.flowsCellprob, tx: t.tx, ty: t.ty, bsize: tileSize }); + tileDiags.push({ tx: t.tx, ty: t.ty, bsize: tileSize, inferenceMs: r.inferenceMs }); + postReply({ + type: 'segment-progress', + reqId, + phase: 'inference', + done: i + 1, + total: tiles.length, + }); + } + + // Average overlapping tile predictions, then run dynamics once. + postReply({ type: 'segment-progress', reqId, phase: 'dynamics', done: 0, total: 1 }); + const tPost = performance.now(); + const averaged = averageTiles(flows, h, w); + const hwFull = h * w; + const dPFull = averaged.data.subarray(0, 2 * hwFull) as Float32Array; + const cpFull = averaged.data.subarray(2 * hwFull, 3 * hwFull) as Float32Array; + const dynOpts: ComputeMasksOptions = { ...(opts.dynamics ?? {}) }; + if (dynOpts.niter === undefined && opts.diameter !== undefined && scale !== 1) { + dynOpts.niter = Math.max(1, Math.floor(200 / scale)); + } + const m = computeMasks(dPFull, cpFull, h, w, dynOpts); + + // Inverse-resize labels back to source resolution (nearest-neighbor). + let masksSrc = m.masks; + let outW = w, + outH = h; + if (scale !== 1) { + outW = image.width; + outH = image.height; + const resized = new Uint32Array(outW * outH); + for (let yy = 0; yy < outH; yy++) { + const sy = Math.min(h - 1, Math.max(0, Math.round(yy * scale))); + const srcRow = sy * w; + const dstRow = yy * outW; + for (let xx = 0; xx < outW; xx++) { + const sx = Math.min(w - 1, Math.max(0, Math.round(xx * scale))); + resized[dstRow + xx] = m.masks[srcRow + sx] as number; + } + } + masksSrc = resized; + } + const postprocessMs = performance.now() - tPost; + postReply( + { + type: 'segment-result', + reqId, + masks: masksSrc, + count: m.count, + width: outW, + height: outH, + resizedWidth: w, + resizedHeight: h, + scale, + tiles: tileDiags, + postprocessMs, + }, + [masksSrc.buffer], + ); } function postReply(msg: WorkerToMain, transfer?: Transferable[]): void { @@ -105,9 +225,7 @@ function postReply(msg: WorkerToMain, transfer?: Transferable[]): void { } // Beacon: confirms the worker module loaded successfully (i.e. the -// `import * as ort` at the top didn't throw). If the user never sees this -// status string, the worker script itself failed to load — most often an -// ORT-web import problem (Firefox without WebGPU, missing WASM sidecars, etc). +// `import * as ort` at the top didn't throw). postReply({ type: 'status', status: 'worker module loaded' }); self.addEventListener('message', async (ev: MessageEvent) => { @@ -117,12 +235,14 @@ self.addEventListener('message', async (ev: MessageEvent) => { postReply({ type: 'status', status: 'init message received' }); await handleInit(msg); } else if (msg.type === 'run-tile') await handleRunTile(msg); + else if (msg.type === 'segment') await handleSegment(msg); else if (msg.type === 'dispose') { session = null; self.close(); } } catch (err) { const message = err instanceof Error ? err.message : String(err); - postReply({ type: 'error', tileId: msg.type === 'run-tile' ? msg.tileId : null, message }); + if (msg.type === 'segment') postReply({ type: 'segment-error', reqId: msg.reqId, message }); + else postReply({ type: 'error', tileId: msg.type === 'run-tile' ? msg.tileId : null, message }); } }); diff --git a/src/worker-protocol.ts b/src/worker-protocol.ts index d441904..9f802e6 100644 --- a/src/worker-protocol.ts +++ b/src/worker-protocol.ts @@ -2,18 +2,56 @@ * Message contract between the public Cellpose API (main thread) and the * inference worker. Discriminated unions on `type` for both directions. * + * Two segmentation paths: + * - `segment` runs the WHOLE pipeline (preprocess → tile → inference → + * average → dynamics → inverse-resize) in the worker and returns the final + * label map. Nothing heavy runs on the main thread, so the UI never blocks. + * - `run-tile` runs a single tile's inference (legacy; the main thread did the + * orchestration). Kept for back-compat. + * * Transferables policy: - * - `MainToWorker.init.modelBytes` is transferred (worker takes ownership). - * - `MainToWorker.runTile.tile` (Float32Array) — the underlying ArrayBuffer - * is transferred so the worker reads it without a copy. - * - `WorkerToMain.tileResult.flowsCellprob` — same: the worker transfers - * the FP32 output back so the main thread can stitch / postprocess - * without copying. + * - `init.modelBytes` is transferred (worker takes ownership). + * - `run-tile.tile` / `tile-result.flowsCellprob` ArrayBuffers are transferred. + * - `segment.image.data` is cloned (left intact for the caller); the returned + * `segment-result.masks` is transferred back (zero-copy). */ +import type { ChannelMapOptions } from './preprocess/channels.js'; +import type { NormalizeOptions } from './preprocess/normalize.js'; +import type { ComputeMasksOptions } from './postprocess/compute_masks.js'; + +/** Serializable subset of SegmentOptions safe to structured-clone to the worker + * (no functions, no AbortSignal). */ +export interface WorkerSegmentOptions extends ChannelMapOptions { + diameter?: number; + tile?: number; + overlap?: number; + normalize?: NormalizeOptions; + dynamics?: ComputeMasksOptions; +} + +/** Per-tile timing diagnostics returned with a full-segment result. The flow + * tensors stay in the worker — returning them would defeat the offload. */ +export interface WorkerTileDiag { + tx: number; + ty: number; + bsize: number; + inferenceMs: number; +} export type MainToWorker = | { type: 'init'; modelBytes: ArrayBuffer; wasmPaths?: string } | { type: 'run-tile'; tileId: number; tile: Float32Array; bsize: number } + | { + type: 'segment'; + reqId: number; + image: { + data: Uint8ClampedArray | Uint8Array | Float32Array; + width: number; + height: number; + channels: number; + }; + opts: WorkerSegmentOptions; + } | { type: 'dispose' }; export type WorkerToMain = @@ -23,4 +61,26 @@ export type WorkerToMain = // Pre-ready progress strings: 'configuring ORT', 'creating session', // 'session created', 'describing adapter'. Optional; consumers wire via // `FromPretrainedOptions.onStatus`. - | { type: 'status'; status: string }; + | { type: 'status'; status: string } + // Full-pipeline segment progress: per-tile inference, then the dynamics phase. + | { + type: 'segment-progress'; + reqId: number; + phase: 'inference' | 'dynamics'; + done: number; + total: number; + } + | { + type: 'segment-result'; + reqId: number; + masks: Uint32Array; + count: number; + width: number; + height: number; + resizedWidth: number; + resizedHeight: number; + scale: number; + tiles: WorkerTileDiag[]; + postprocessMs: number; + } + | { type: 'segment-error'; reqId: number; message: string };