From 23c1b3e6967c82c6e1fb8bd56e40b5e0bf526f39 Mon Sep 17 00:00:00 2001 From: belkassaby Date: Fri, 22 May 2026 21:04:31 -0400 Subject: [PATCH 1/5] fix python parity in preprocess + dynamics; expand test coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five corrections after a code review against cellpose (Python): 1. Preprocess order: normalize FIRST, then resize. Matches `models.py:_run_net`. Resizing before normalize was changing the per-channel percentile statistics. 2. Replace canvas-based resize with pure-JS Float32 bilinear. The old path quantized each channel to uint8 (~1/255 error per pixel) and required OffscreenCanvas/HTMLCanvas. New impl uses OpenCV's INTER_LINEAR pixel-center mapping with edge replication. 3. Scale niter by image_scaling. Python computes `niter = int(200 / (30 / diameter))`; JS was fixed at 200 regardless of diameter. Explicit user override still wins. 4. Math.round → Math.trunc in follow_flows. PyTorch's `.int()` cast truncates toward zero; JS was rounding-half-up, which shifted converged points by up to 1 pixel. 5. Fix README quickstart: `cellprob_threshold` was at the wrong nesting (should be `dynamics: { cellprobThreshold }`); also removed a stale "once published" parenthetical. Quality impact: - `three_cells_192` dynamics parity: 0.601 → 1.000 mean IoU (all 5 per-cell IoUs at 1.000). Gate tightened from 0.55 → 0.99. - Total tests 14 → 61. New files cover `buildCpsamChannels`, `percentile`, `taperMask`, `followFlows`, `getMasks`, the new bilinear `diameterResize`, and an end-to-end postprocess pipeline test that mocks the inference step. Co-Authored-By: Claude Opus 4.7 (1M context) --- README.md | 4 +- src/cellpose.ts | 16 ++- src/postprocess/follow_flows.ts | 8 +- src/preprocess/resize.ts | 128 ++++++++------------- tests/channels.test.ts | 75 ++++++++++++ tests/dynamics.test.ts | 13 +-- tests/end_to_end.test.ts | 194 ++++++++++++++++++++++++++++++++ tests/follow_flows.test.ts | 158 ++++++++++++++++++++++++++ tests/get_masks.test.ts | 157 ++++++++++++++++++++++++++ tests/percentile.test.ts | 62 ++++++++++ tests/resize.test.ts | 126 +++++++++++++++++++++ tests/resize_guard.test.ts | 58 ++++++++++ tests/taper_mask.test.ts | 99 ++++++++++++++++ 13 files changed, 1006 insertions(+), 92 deletions(-) create mode 100644 tests/channels.test.ts create mode 100644 tests/end_to_end.test.ts create mode 100644 tests/follow_flows.test.ts create mode 100644 tests/get_masks.test.ts create mode 100644 tests/percentile.test.ts create mode 100644 tests/resize.test.ts create mode 100644 tests/resize_guard.test.ts create mode 100644 tests/taper_mask.test.ts diff --git a/README.md b/README.md index 5a91e5c..89600c5 100644 --- a/README.md +++ b/README.md @@ -34,7 +34,7 @@ npm install cellpose-js onnxruntime-web You also need to host: -1. **The model**: `cpsam_fp16.onnx` (588 MB). Either upload to your own CDN, or use the public copy at `https://huggingface.co/ballon999/cellpose-sam-onnx/resolve/main/cpsam_fp16.onnx` (once published in M6 follow-up). +1. **The model**: `cpsam_fp16.onnx` (588 MB). Either upload to your own CDN, or use the public copy at `https://huggingface.co/ballon999/cellpose-sam-onnx/resolve/main/cpsam_fp16.onnx`. 2. **ORT-web's WASM/JSEP sidecars**: ORT dynamically imports `.mjs` and `.wasm` files at runtime. They must be served **same-origin** with your app (cross-origin dynamic `import()` is blocked). Either copy `node_modules/onnxruntime-web/dist/ort-wasm-simd-threaded.{wasm,mjs,jsep.wasm,jsep.mjs,asyncify.wasm,asyncify.mjs}` to your public assets, or proxy `/ort/*` to jsDelivr at build time (see `examples/demo/vite.config.ts` for the recipe). ## Quickstart @@ -60,9 +60,9 @@ const result = await cp.segment( { data: imageData.data, width: imageData.width, height: imageData.height, channels: 4 }, { diameter: 30, // estimated cell diameter in source pixels (omit for native resolution) - cellprob_threshold: 0, chan: 0, // primary channel (0 = grayscale) chan2: 0, // secondary channel (0 = none) + dynamics: { cellprobThreshold: 0 }, // pixels above this enter the dynamical system onTileProgress: (done, total) => console.log(`tile ${done}/${total}`), }, ); diff --git a/src/cellpose.ts b/src/cellpose.ts index 97c59bc..de551d2 100644 --- a/src/cellpose.ts +++ b/src/cellpose.ts @@ -249,6 +249,11 @@ export class Cellpose { 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; @@ -256,7 +261,6 @@ export class Cellpose { h = r.height; scale = r.scale; } - chw = normalizePerChannel(chw, 3, w * h, opts.normalize ?? {}); const tileOpts: { bsize: number; overlap?: number } = { bsize: tileSize }; if (opts.overlap !== undefined) tileOpts.overlap = opts.overlap; @@ -288,7 +292,15 @@ export class Cellpose { 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 m = computeMasks(dPFull, cpFull, h, w, opts.dynamics ?? {}); + // 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); // Inverse-resize labels back to source resolution if a diameter resize // was applied (nearest-neighbor — labels must not be interpolated). diff --git a/src/postprocess/follow_flows.ts b/src/postprocess/follow_flows.ts index a188c3e..f4a4553 100644 --- a/src/postprocess/follow_flows.ts +++ b/src/postprocess/follow_flows.ts @@ -131,11 +131,13 @@ export function followFlows( } } - // Un-normalize and round to int. + // Un-normalize and truncate to int. PyTorch's `.int()` cast truncates toward + // zero (it's a C-style cast), not round-to-nearest. gY/gX are clamped to + // [-1, 1] so (g+1)*0.5*(N-1) ∈ [0, N-1] and Math.trunc == Math.floor here. const pFinal = new Int32Array(2 * n); for (let i = 0; i < n; i++) { - pFinal[2 * i] = Math.round(((gY[i] as number) + 1) * 0.5 * (H - 1)); - pFinal[2 * i + 1] = Math.round(((gX[i] as number) + 1) * 0.5 * (W - 1)); + pFinal[2 * i] = Math.trunc(((gY[i] as number) + 1) * 0.5 * (H - 1)); + pFinal[2 * i + 1] = Math.trunc(((gX[i] as number) + 1) * 0.5 * (W - 1)); } return { pFinal, seedY, seedX }; } diff --git a/src/preprocess/resize.ts b/src/preprocess/resize.ts index b501987..1400843 100644 --- a/src/preprocess/resize.ts +++ b/src/preprocess/resize.ts @@ -2,15 +2,15 @@ * Diameter-aware resize, mirroring cellpose.transforms.resize_image with * cv2.INTER_LINEAR (bilinear) as the interpolation. * - * Browser bilinear resize is not bit-exact to OpenCV's — they use different - * edge handling and rounding. We expect mean abs error ~1e-3 vs cv2; the - * fixture-based parity tests use a wider tolerance for resize than for - * pure-math ops like normalize. + * Pure-JS Float32 bilinear — no canvas roundtrip. Earlier versions of this + * module pushed each channel through an OffscreenCanvas which quantized + * values to uint8 (~1/255 relative error per pixel). The canvas path also + * placed an environment dependency on OffscreenCanvas / HTMLCanvas, which + * isn't available in Node/test contexts. * - * Implementation: use OffscreenCanvas where available (workers), HTMLCanvas - * fallback for main-thread contexts. We resize one channel at a time as a - * grayscale image (R = value, A = 255), read back as RGBA, and pull the R - * channel. + * Pixel-center mapping matches OpenCV's INTER_LINEAR: + * src_y = (dst_y + 0.5) * (srcH / dstH) - 0.5 + * with edge replication (BORDER_REPLICATE) outside [0, N-1]. */ export interface ResizeResult { @@ -33,41 +33,12 @@ export interface DiameterResizeOptions { targetDiameter?: number; } -/** Create a 2D canvas context. Prefers OffscreenCanvas for worker contexts. */ -function createCanvas( - w: number, - h: number, -): { - ctx: OffscreenCanvasRenderingContext2D | CanvasRenderingContext2D; -} { - if (typeof OffscreenCanvas !== 'undefined') { - const c = new OffscreenCanvas(w, h); - const ctx = c.getContext('2d'); - if (!ctx) throw new Error('OffscreenCanvas 2d context unavailable'); - return { ctx }; - } - if (typeof document !== 'undefined') { - const c = document.createElement('canvas'); - c.width = w; - c.height = h; - const ctx = c.getContext('2d'); - if (!ctx) throw new Error('HTMLCanvas 2d context unavailable'); - return { ctx }; - } - throw new Error('No canvas context available in this environment'); -} - /** - * Resize one channel using canvas bilinear. - * - * The channel data is packed into RGBA (R = value scaled to [0,255]), drawn, - * resized, read back as RGBA, and the R channel is rescaled to the original - * value range. + * Resize one channel with bilinear interpolation on Float32 values. * - * This is lossy — values are quantized to uint8 in the canvas pipeline. For - * percentile-normalized inputs in roughly [0, 1] this is acceptable (~1/255 - * quantization error). For raw images use the source-value range; callers - * should normalize *after* resize if they need full precision. + * Pixel-center mapping is OpenCV's INTER_LINEAR convention: + * src_y = (dst_y + 0.5) * (srcH / dstH) - 0.5 + * Out-of-bounds samples replicate the nearest edge pixel (BORDER_REPLICATE). */ function resizeChannel( src: Float32Array, @@ -76,41 +47,42 @@ function resizeChannel( dstW: number, dstH: number, ): Float32Array { - // Find the channel's value range so quantization uses the full uint8. - let mn = Infinity, - mx = -Infinity; - for (let i = 0; i < src.length; i++) { - const v = src[i] as number; - if (v < mn) mn = v; - if (v > mx) mx = v; - } - const span = mx - mn || 1; - - // Build the source ImageData. - const srcRgba = new Uint8ClampedArray(srcW * srcH * 4); - for (let i = 0; i < src.length; i++) { - const u8 = Math.max(0, Math.min(255, Math.round((((src[i] as number) - mn) / span) * 255))); - srcRgba[i * 4] = u8; - srcRgba[i * 4 + 1] = u8; - srcRgba[i * 4 + 2] = u8; - srcRgba[i * 4 + 3] = 255; - } - const srcImageData = new ImageData(srcRgba, srcW, srcH); - - const { ctx: srcCtx } = createCanvas(srcW, srcH); - srcCtx.putImageData(srcImageData, 0, 0); - - const { ctx: dstCtx } = createCanvas(dstW, dstH); - dstCtx.imageSmoothingEnabled = true; - dstCtx.imageSmoothingQuality = 'high'; - // drawImage with a canvas source supports floating-point resizing via bilinear. - // OffscreenCanvas accepts itself as a CanvasImageSource. - dstCtx.drawImage(srcCtx.canvas as CanvasImageSource, 0, 0, dstW, dstH); - - const dstRgba = dstCtx.getImageData(0, 0, dstW, dstH).data; const out = new Float32Array(dstW * dstH); - for (let i = 0; i < out.length; i++) { - out[i] = ((dstRgba[i * 4] as number) / 255) * span + mn; + const syScale = srcH / dstH; + const sxScale = srcW / dstW; + for (let dy = 0; dy < dstH; dy++) { + const sy = (dy + 0.5) * syScale - 0.5; + let y0 = Math.floor(sy); + let y1 = y0 + 1; + const fy = sy - y0; + // Edge-replicate clamps. Hot-path branchless mins/maxes are not worth the + // readability hit here — modern engines hoist these. + if (y0 < 0) y0 = 0; + else if (y0 > srcH - 1) y0 = srcH - 1; + if (y1 < 0) y1 = 0; + else if (y1 > srcH - 1) y1 = srcH - 1; + const row0 = y0 * srcW; + const row1 = y1 * srcW; + const dstRow = dy * dstW; + const wy1 = fy; + const wy0 = 1 - fy; + for (let dx = 0; dx < dstW; dx++) { + const sx = (dx + 0.5) * sxScale - 0.5; + let x0 = Math.floor(sx); + let x1 = x0 + 1; + const fx = sx - x0; + if (x0 < 0) x0 = 0; + else if (x0 > srcW - 1) x0 = srcW - 1; + if (x1 < 0) x1 = 0; + else if (x1 > srcW - 1) x1 = srcW - 1; + const v00 = src[row0 + x0] as number; + const v01 = src[row0 + x1] as number; + const v10 = src[row1 + x0] as number; + const v11 = src[row1 + x1] as number; + const v0 = v00 * (1 - fx) + v01 * fx; + const v1 = v10 * (1 - fx) + v11 * fx; + out[dstRow + dx] = v0 * wy0 + v1 * wy1; + } } return out; } @@ -140,9 +112,9 @@ export function diameterResize( } const dstW = Math.max(1, Math.round(width * scale)); const dstH = Math.max(1, Math.round(height * scale)); - // Memory guard: the canvas-based per-channel resize allocates ~3 canvases - // per channel × 3 channels. A 4096x4096 destination already consumes - // ~300 MB across those canvases. Anything larger reliably OOMs. + // Memory guard: each resized channel allocates dstW*dstH Float32 values. + // For a 4096×4096 destination at 3 channels that's ~192 MB; anything + // larger reliably OOMs on memory-constrained devices. const MAX_PIXELS = 4096 * 4096; if (dstW * dstH > MAX_PIXELS) { throw new Error( diff --git a/tests/channels.test.ts b/tests/channels.test.ts new file mode 100644 index 0000000..a3ef9f7 --- /dev/null +++ b/tests/channels.test.ts @@ -0,0 +1,75 @@ +/** + * Unit tests for buildCpsamChannels (preprocess/channels.ts). + * + * Coverage: + * - chan=0 grayscale: single source channel replicates into output 0. + * - chan=1/2/3 picks R/G/B from a pixel-interleaved source. + * - chan2 != 0 places the secondary channel in output slot 1. + * - Output slot 2 is always zero. + * - Layout: output is CHW (channel-major), so adjacent floats within a + * channel are pixel neighbors. + * - Error paths: wrong total length; out-of-range channel index. + */ +import { describe, it, expect } from 'vitest'; +import { buildCpsamChannels } from '../src/preprocess/channels.js'; + +describe('buildCpsamChannels', () => { + it('grayscale source (channels=1, chan=0) replicates to output 0', () => { + // 2x2 grayscale: pixel values 10, 20, 30, 40. + const src = new Uint8Array([10, 20, 30, 40]); + const out = buildCpsamChannels(src, 2, 2, 1, { chan: 0 }); + // (3, 2, 2) CHW: channel 0 = grayscale, channels 1/2 = zeros. + expect(out.length).toBe(3 * 2 * 2); + expect(Array.from(out.slice(0, 4))).toEqual([10, 20, 30, 40]); // channel 0 + expect(Array.from(out.slice(4, 8))).toEqual([0, 0, 0, 0]); // channel 1 + expect(Array.from(out.slice(8, 12))).toEqual([0, 0, 0, 0]); // channel 2 + }); + + it('RGBA source (channels=4): chan=1 picks R, chan2=3 picks B', () => { + // 2x1 RGBA: pixel0 = (R=10, G=11, B=12, A=255), pixel1 = (R=20, G=21, B=22, A=255). + const src = new Uint8Array([10, 11, 12, 255, 20, 21, 22, 255]); + const out = buildCpsamChannels(src, 2, 1, 4, { chan: 1, chan2: 3 }); + // Channel 0 = R = [10, 20] + // Channel 1 = B = [12, 22] + // Channel 2 = zeros + expect(Array.from(out.slice(0, 2))).toEqual([10, 20]); + expect(Array.from(out.slice(2, 4))).toEqual([12, 22]); + expect(Array.from(out.slice(4, 6))).toEqual([0, 0]); + }); + + it('RGB source (channels=3): chan=2 picks G, chan2=0 leaves slot 1 empty', () => { + const src = new Uint8Array([1, 2, 3, 4, 5, 6]); // 2 pixels RGB + const out = buildCpsamChannels(src, 2, 1, 3, { chan: 2, chan2: 0 }); + // Channel 0 = G = [2, 5] + expect(Array.from(out.slice(0, 2))).toEqual([2, 5]); + expect(Array.from(out.slice(2, 4))).toEqual([0, 0]); + expect(Array.from(out.slice(4, 6))).toEqual([0, 0]); + }); + + it('Float32 source preserves precision (no quantization)', () => { + const src = new Float32Array([0.1, 0.2, 0.3, 0.4]); + const out = buildCpsamChannels(src, 2, 2, 1, { chan: 0 }); + expect(out[0]).toBeCloseTo(0.1, 6); + expect(out[3]).toBeCloseTo(0.4, 6); + }); + + it('defaults to chan=0, chan2=0 (grayscale-equivalent)', () => { + const src = new Uint8Array([10, 20, 30, 40]); + const out = buildCpsamChannels(src, 2, 2, 1); + expect(Array.from(out.slice(0, 4))).toEqual([10, 20, 30, 40]); + expect(Array.from(out.slice(4, 8))).toEqual([0, 0, 0, 0]); + }); + + it('throws on length mismatch', () => { + const src = new Uint8Array([1, 2, 3]); // 3 floats for a 2x2 grayscale (needs 4) + expect(() => buildCpsamChannels(src, 2, 2, 1)).toThrow(/expected 4/); + }); + + it('throws on channel index exceeding source channel count', () => { + const src = new Uint8Array([1, 2, 3, 4]); // 2x2 grayscale + // chan=2 requires a source with >=2 channels. + expect(() => buildCpsamChannels(src, 2, 2, 1, { chan: 2 })).toThrow( + /channel index 2 not available/, + ); + }); +}); diff --git a/tests/dynamics.test.ts b/tests/dynamics.test.ts index 40ec513..a5cd352 100644 --- a/tests/dynamics.test.ts +++ b/tests/dynamics.test.ts @@ -26,13 +26,12 @@ interface Case { describe('compute_masks parity', () => { const cases: Case[] = [ { name: 'single_cell_128', gate: 0.9 }, - // three_cells_192's synthetic flow fields have overlapping basins of - // attraction where JS and Python converge to slightly different seed - // structures. 3 of 5 Python labels are matched at IoU > 0.95; 2 small - // labels in the overlap zone have IoU ≈ 0 because they're in - // numerical-noise territory. Gate set accordingly; the real M4 gate - // (mean IoU ≥ 0.9 on real CPSAM outputs) lives in a follow-up fixture. - { name: 'three_cells_192', gate: 0.55 }, + // Pre-fix this fixture sat at 0.601 because JS used Math.round on the + // final convergence coordinates while PyTorch's `.int()` truncates. The + // 1-pixel shift split two small overlap-zone labels off their Python + // counterparts. After switching to Math.trunc all 5 labels match at IoU + // 1.000, so the gate is tightened here to catch any regression. + { name: 'three_cells_192', gate: 0.99 }, { name: 'empty_96', gate: 0.99 }, ]; for (const c of cases) { diff --git a/tests/end_to_end.test.ts b/tests/end_to_end.test.ts new file mode 100644 index 0000000..bda7700 --- /dev/null +++ b/tests/end_to_end.test.ts @@ -0,0 +1,194 @@ +/** + * End-to-end test for the postprocess pipeline (tile → average → dynamics). + * + * The inference step itself runs on WebGPU and isn't reachable from the + * Node test runner, so we mock the per-tile model output: we synthesize a + * flow field that points every pixel toward a known per-cell center, plus + * a cellprob > threshold for those pixels. The pipeline should recover the + * cells we planted. + * + * This catches integration bugs across averageTiles + computeMasks that + * the module-level parity tests miss (e.g. wrong channel-order assumptions, + * tile-coordinate vs full-image-coordinate mistakes). + */ +import { describe, it, expect } from 'vitest'; +import { makeTiles } from '../src/preprocess/tile.js'; +import { averageTiles, type TileInputForAveraging } from '../src/postprocess/average_tiles.js'; +import { computeMasks } from '../src/postprocess/compute_masks.js'; +import { instanceIoUs } from './util/iou.js'; + +interface Cell { + cy: number; + cx: number; + radius: number; +} + +/** + * Synthesize a full-image (3, H, W) flow+cellprob tensor where every pixel + * inside any cell's radius points toward its center, and cellprob inside + * cells is 1.0 (0 elsewhere). + * + * Layout: [dy(0..hw), dx(hw..2hw), cellprob(2hw..3hw)] in CHW row-major. + */ +function synthesizeFullImageFlows(H: number, W: number, cells: Cell[]): Float32Array { + const hw = H * W; + const out = new Float32Array(3 * hw); + for (let y = 0; y < H; y++) { + for (let x = 0; x < W; x++) { + const i = y * W + x; + let bestCell: Cell | null = null; + let bestD = Infinity; + for (const c of cells) { + const d = Math.hypot(y - c.cy, x - c.cx); + if (d <= c.radius && d < bestD) { + bestCell = c; + bestD = d; + } + } + if (bestCell !== null) { + const dy = bestCell.cy - y; + const dx = bestCell.cx - x; + const r = Math.hypot(dy, dx) || 1; + out[i] = dy / r; // unit vector toward center; compute_masks scales by 1/5 + out[hw + i] = dx / r; + out[2 * hw + i] = 1.0; // cellprob + } + } + } + return out; +} + +/** + * Slice the full-image (3, H, W) tensor into per-tile (3, B, B) tensors. + * This mirrors what real inference produces: per-tile outputs in tile-local + * coordinates, ready for averageTiles to stitch back together. + */ +function sliceFullImageToTiles( + full: Float32Array, + H: number, + W: number, + bsize: number, + tileOrigins: Array<{ tx: number; ty: number }>, +): TileInputForAveraging[] { + const hw = H * W; + const tileHW = bsize * bsize; + const out: TileInputForAveraging[] = []; + for (const origin of tileOrigins) { + const flowsCellprob = new Float32Array(3 * tileHW); + for (let c = 0; c < 3; c++) { + for (let dy = 0; dy < bsize; dy++) { + const yy = origin.ty + dy; + if (yy < 0 || yy >= H) continue; + for (let dx = 0; dx < bsize; dx++) { + const xx = origin.tx + dx; + if (xx < 0 || xx >= W) continue; + flowsCellprob[c * tileHW + dy * bsize + dx] = full[c * hw + yy * W + xx] as number; + } + } + } + out.push({ flowsCellprob, tx: origin.tx, ty: origin.ty, bsize }); + } + return out; +} + +/** + * Reference ground-truth label map: every pixel within a cell's radius gets + * that cell's 1-indexed label. + */ +function buildGtLabels(H: number, W: number, cells: Cell[]): Uint32Array { + const out = new Uint32Array(H * W); + for (let y = 0; y < H; y++) { + for (let x = 0; x < W; x++) { + for (let k = 0; k < cells.length; k++) { + const c = cells[k]!; + if (Math.hypot(y - c.cy, x - c.cx) <= c.radius) { + out[y * W + x] = k + 1; + break; + } + } + } + } + return out; +} + +describe('postprocess pipeline (averageTiles + computeMasks)', () => { + it('single tile, 1 synthetic cell → 1 label that covers the cell', () => { + const H = 96, + W = 96; + const cells: Cell[] = [{ cy: 48, cx: 48, radius: 12 }]; + const full = synthesizeFullImageFlows(H, W, cells); + // No tiling — pass a single full-image "tile". + const tiles: TileInputForAveraging[] = [ + { flowsCellprob: full, tx: 0, ty: 0, bsize: Math.max(H, W) }, + ]; + // averageTiles expects bsize-sized inputs, so reshape full → (3, bsize, bsize) + // For this test we just call it with bsize=H=W (square). + const avg = averageTiles(tiles, H, W); + const hwFull = H * W; + const dP = avg.data.subarray(0, 2 * hwFull) as Float32Array; + const cp = avg.data.subarray(2 * hwFull, 3 * hwFull) as Float32Array; + const r = computeMasks(dP, cp, H, W, { niter: 300 }); + expect(r.count).toBe(1); + const gt = buildGtLabels(H, W, cells); + const { mean } = instanceIoUs(gt, r.masks); + expect(mean).toBeGreaterThan(0.85); + }); + + it('multi-tile 2x2 grid, 4 cells (one per quadrant) → 4 labels', () => { + const H = 400, + W = 400; + const bsize = 256; + const cells: Cell[] = [ + { cy: 100, cx: 100, radius: 18 }, + { cy: 100, cx: 300, radius: 18 }, + { cy: 300, cx: 100, radius: 18 }, + { cy: 300, cx: 300, radius: 18 }, + ]; + const full = synthesizeFullImageFlows(H, W, cells); + // Use makeTiles on a dummy image just to get the tile origins. + const dummy = new Float32Array(3 * H * W); + const tileRecs = makeTiles(dummy, W, H, 3, { bsize, overlap: 0.1 }); + const tileInputs = sliceFullImageToTiles( + full, + H, + W, + bsize, + tileRecs.map((t) => ({ tx: t.tx, ty: t.ty })), + ); + const avg = averageTiles(tileInputs, H, W); + const hwFull = H * W; + const dP = avg.data.subarray(0, 2 * hwFull) as Float32Array; + const cp = avg.data.subarray(2 * hwFull, 3 * hwFull) as Float32Array; + const r = computeMasks(dP, cp, H, W, { niter: 300 }); + expect(r.count).toBe(4); + const gt = buildGtLabels(H, W, cells); + const { mean, per } = instanceIoUs(gt, r.masks); + // Per-cell IoU should be high (each cell is well-isolated). + expect(per.length).toBe(4); + for (const iou of per) expect(iou).toBeGreaterThan(0.7); + expect(mean).toBeGreaterThan(0.8); + }); + + it('cellprobThreshold filters out a cell when raised above its cellprob', () => { + const H = 96, + W = 96; + const cells: Cell[] = [{ cy: 48, cx: 48, radius: 12 }]; + const full = synthesizeFullImageFlows(H, W, cells); + const hwFull = H * W; + const dP = full.subarray(0, 2 * hwFull) as Float32Array; + const cp = full.subarray(2 * hwFull, 3 * hwFull) as Float32Array; + // Cellprob is 1.0 inside the cell. Set threshold to 1.5 → no pixel passes. + const r = computeMasks(dP, cp, H, W, { cellprobThreshold: 1.5 }); + expect(r.count).toBe(0); + for (let i = 0; i < r.masks.length; i++) expect(r.masks[i]).toBe(0); + }); + + it('empty cellprob → zero masks (defensive path)', () => { + const H = 64, + W = 64; + const dP = new Float32Array(2 * H * W); + const cp = new Float32Array(H * W); // all zero + const r = computeMasks(dP, cp, H, W); + expect(r.count).toBe(0); + }); +}); diff --git a/tests/follow_flows.test.ts b/tests/follow_flows.test.ts new file mode 100644 index 0000000..6cf35a7 --- /dev/null +++ b/tests/follow_flows.test.ts @@ -0,0 +1,158 @@ +/** + * Unit tests for followFlows (postprocess/follow_flows.ts). + * + * Strategy: synthesize a flow field with a known fixed point and verify that + * Euler integration drives seed pixels toward it. This is the lightest-weight + * way to anchor the Euler loop without depending on a real CPSAM output. + * + * NOTE: cellpose pre-scales the flow field by `(cellprob > thresh) / 5` + * before calling followFlows. These tests pass already-scaled inputs so + * each step's displacement is the desired fraction of a pixel. + */ +import { describe, it, expect } from 'vitest'; +import { followFlows } from '../src/postprocess/follow_flows.js'; + +/** + * Build a flow field where every pixel's (dy, dx) points toward (cy, cx). + * Pre-scales by 1/5 to match the convention used by compute_masks. + */ +function flowsTowardCenter( + H: number, + W: number, + cy: number, + cx: number, + magnitude = 1.0, +): { dP: Float32Array; cellprob: Float32Array } { + const hw = H * W; + const dP = new Float32Array(2 * hw); + const cellprob = new Float32Array(hw); + for (let y = 0; y < H; y++) { + for (let x = 0; x < W; x++) { + const i = y * W + x; + const dy = cy - y; + const dx = cx - x; + const r = Math.hypot(dy, dx) || 1; + // Unit vector toward center, scaled by magnitude, then by 1/5 + // (the compute_masks convention). + dP[i] = ((dy / r) * magnitude) / 5; + dP[hw + i] = ((dx / r) * magnitude) / 5; + cellprob[i] = 1.0; // every pixel is "cell" so every pixel is a seed + } + } + return { dP, cellprob }; +} + +describe('followFlows', () => { + it('seeds with cellprob <= threshold are excluded', () => { + const H = 16, + W = 16; + const { dP, cellprob } = flowsTowardCenter(H, W, 8, 8); + // Knock out one row of cellprob so those seeds drop out. + for (let x = 0; x < W; x++) cellprob[3 * W + x] = 0; + const r = followFlows(dP, cellprob, H, W, 0, 50); + expect(r.seedY.length).toBe(H * W - W); + // None of the surviving seeds should come from y=3. + for (let i = 0; i < r.seedY.length; i++) { + expect(r.seedY[i]).not.toBe(3); + } + }); + + it('returns empty arrays when no pixel exceeds the threshold', () => { + const H = 8, + W = 8; + const dP = new Float32Array(2 * H * W); + const cellprob = new Float32Array(H * W); // all zero + const r = followFlows(dP, cellprob, H, W, 0, 10); + expect(r.pFinal.length).toBe(0); + expect(r.seedY.length).toBe(0); + expect(r.seedX.length).toBe(0); + }); + + it('converges to the fixed point under sufficient iterations', () => { + const H = 32, + W = 32; + const cy = 15, + cx = 16; + const { dP, cellprob } = flowsTowardCenter(H, W, cy, cx, 1.0); + const r = followFlows(dP, cellprob, H, W, 0, 300); + expect(r.seedY.length).toBe(H * W); + // Every seed should converge to within a few pixels of (cy, cx). + // The fixed point is a sink — once a pixel arrives, the flow magnitude + // at the fixed point itself is undefined (we set unit magnitude / 5), + // so we accept any final position within a small radius. + let maxDist = 0; + for (let i = 0; i < r.pFinal.length / 2; i++) { + const py = r.pFinal[2 * i] as number; + const px = r.pFinal[2 * i + 1] as number; + const d = Math.hypot(py - cy, px - cx); + if (d > maxDist) maxDist = d; + } + expect(maxDist).toBeLessThan(4); + }); + + it('two basins of attraction → seeds converge to nearest center', () => { + const H = 48, + W = 48; + const c1 = { y: 12, x: 12 }; + const c2 = { y: 36, x: 36 }; + const hw = H * W; + const dP = new Float32Array(2 * hw); + const cellprob = new Float32Array(hw); + for (let y = 0; y < H; y++) { + for (let x = 0; x < W; x++) { + const i = y * W + x; + const d1 = Math.hypot(y - c1.y, x - c1.x); + const d2 = Math.hypot(y - c2.y, x - c2.x); + const c = d1 < d2 ? c1 : c2; + const dy = c.y - y; + const dx = c.x - x; + const r = Math.hypot(dy, dx) || 1; + dP[i] = (dy / r) / 5; + dP[hw + i] = (dx / r) / 5; + cellprob[i] = 1.0; + } + } + const r = followFlows(dP, cellprob, H, W, 0, 400); + let nearC1 = 0, + nearC2 = 0; + for (let i = 0; i < r.pFinal.length / 2; i++) { + const py = r.pFinal[2 * i] as number; + const px = r.pFinal[2 * i + 1] as number; + const d1 = Math.hypot(py - c1.y, px - c1.x); + const d2 = Math.hypot(py - c2.y, px - c2.x); + if (d1 < 5) nearC1++; + else if (d2 < 5) nearC2++; + } + // Both clusters should accumulate roughly equal seeds (we sized each + // basin symmetrically). The exact split depends on the diagonal boundary, + // but both must capture a sizeable share. + expect(nearC1).toBeGreaterThan(H * W * 0.3); + expect(nearC2).toBeGreaterThan(H * W * 0.3); + expect(nearC1 + nearC2).toBeGreaterThan(H * W * 0.95); + }); + + it('final coordinates are bounded by [0, H-1] x [0, W-1]', () => { + const H = 24, + W = 24; + // Strong flow that would push past the boundary if not clamped. + const { dP, cellprob } = flowsTowardCenter(H, W, 12, 12, 100.0); + const r = followFlows(dP, cellprob, H, W, 0, 200); + for (let i = 0; i < r.pFinal.length / 2; i++) { + const py = r.pFinal[2 * i] as number; + const px = r.pFinal[2 * i + 1] as number; + expect(py).toBeGreaterThanOrEqual(0); + expect(py).toBeLessThanOrEqual(H - 1); + expect(px).toBeGreaterThanOrEqual(0); + expect(px).toBeLessThanOrEqual(W - 1); + } + }); + + it('throws on flow / cellprob size mismatch', () => { + const H = 8, + W = 8; + const goodDP = new Float32Array(2 * H * W); + const goodCP = new Float32Array(H * W); + expect(() => followFlows(new Float32Array(7), goodCP, H, W)).toThrow(/dP wrong size/); + expect(() => followFlows(goodDP, new Float32Array(7), H, W)).toThrow(/cellprob wrong size/); + }); +}); diff --git a/tests/get_masks.test.ts b/tests/get_masks.test.ts new file mode 100644 index 0000000..ec83118 --- /dev/null +++ b/tests/get_masks.test.ts @@ -0,0 +1,157 @@ +/** + * Unit tests for getMasks (postprocess/get_masks.ts). + * + * Strategy: hand-build a "post-Euler" convergence pattern with N tight + * clusters of identical (py, px) points. Each cluster should produce one + * mask covering its source seed pixels. + */ +import { describe, it, expect } from 'vitest'; +import { getMasks } from '../src/postprocess/get_masks.js'; + +/** + * Build N synthetic clusters: M seeds per cluster, all converging to one + * point. Source seed pixels for cluster k are arranged in a small square + * starting at (sourceCenters[k].y, sourceCenters[k].x). + */ +function makeClusters( + H: number, + W: number, + clusters: Array<{ converge: { y: number; x: number }; source: { y: number; x: number } }>, + perCluster: number, +): { + pFinal: Int32Array; + seedY: Int32Array; + seedX: Int32Array; +} { + const n = clusters.length * perCluster; + const pFinal = new Int32Array(2 * n); + const seedY = new Int32Array(n); + const seedX = new Int32Array(n); + let i = 0; + for (const c of clusters) { + // Arrange perCluster seeds in a roughly square footprint at source. + const side = Math.ceil(Math.sqrt(perCluster)); + for (let k = 0; k < perCluster; k++) { + const sy = c.source.y + Math.floor(k / side); + const sx = c.source.x + (k % side); + // Wrap inside the image to keep the test self-consistent. + seedY[i] = Math.min(H - 1, sy); + seedX[i] = Math.min(W - 1, sx); + pFinal[2 * i] = c.converge.y; + pFinal[2 * i + 1] = c.converge.x; + i++; + } + } + return { pFinal, seedY, seedX }; +} + +describe('getMasks', () => { + it('returns zero-count for empty input', () => { + const r = getMasks(new Int32Array(0), new Int32Array(0), new Int32Array(0), 64, 64); + expect(r.count).toBe(0); + expect(r.masks.length).toBe(64 * 64); + for (let i = 0; i < r.masks.length; i++) expect(r.masks[i]).toBe(0); + }); + + it('produces one label per cluster (3 clusters)', () => { + const H = 96, + W = 96; + const clusters = [ + { converge: { y: 16, x: 16 }, source: { y: 8, x: 8 } }, + { converge: { y: 48, x: 48 }, source: { y: 40, x: 40 } }, + { converge: { y: 80, x: 16 }, source: { y: 72, x: 8 } }, + ]; + // ~25 seeds per cluster — well above the PEAK_HIST_THRESHOLD of 10. + const { pFinal, seedY, seedX } = makeClusters(H, W, clusters, 25); + const r = getMasks(pFinal, seedY, seedX, H, W); + expect(r.count).toBe(3); + const seen = new Set(); + for (let i = 0; i < r.masks.length; i++) { + const v = r.masks[i] as number; + if (v > 0) seen.add(v); + } + expect(seen.size).toBe(3); + expect(Array.from(seen).sort((a, b) => a - b)).toEqual([1, 2, 3]); + }); + + it('drops clusters that fall under PEAK_HIST_THRESHOLD (10)', () => { + const H = 64, + W = 64; + const clusters = [ + { converge: { y: 16, x: 16 }, source: { y: 8, x: 8 } }, // 5 seeds → dropped + { converge: { y: 48, x: 48 }, source: { y: 40, x: 40 } }, // 25 seeds → kept + ]; + // perCluster = 5 → both clusters get 5 seeds. The first cluster's peak in + // the histogram is therefore 5, which is below the threshold (10), so it + // should be dropped. + const { pFinal, seedY, seedX } = makeClusters(H, W, clusters, 5); + const r1 = getMasks(pFinal, seedY, seedX, H, W); + expect(r1.count).toBe(0); // both clusters below threshold + + const { pFinal: p2, seedY: y2, seedX: x2 } = makeClusters(H, W, clusters, 25); + const r2 = getMasks(p2, y2, x2, H, W); + expect(r2.count).toBe(2); + }); + + it('drops masks larger than maxSizeFraction of the image', () => { + const H = 64, + W = 64; + const total = H * W; + // One giant cluster: 60% of pixels (above the default 40% threshold). + // We need at least PEAK_HIST_THRESHOLD+1 (=11) seeds piling onto the + // convergence point for it to be detected as a seed in the first place. + const targetSize = Math.floor(total * 0.6); + const pFinal = new Int32Array(2 * targetSize); + const seedY = new Int32Array(targetSize); + const seedX = new Int32Array(targetSize); + for (let i = 0; i < targetSize; i++) { + seedY[i] = Math.floor(i / W); + seedX[i] = i % W; + pFinal[2 * i] = 32; + pFinal[2 * i + 1] = 32; + } + const r = getMasks(pFinal, seedY, seedX, H, W, 0.4); + // Mask exceeds 40% → dropped → 0 labels remain. + expect(r.count).toBe(0); + for (let i = 0; i < r.masks.length; i++) expect(r.masks[i]).toBe(0); + }); + + it('renumbers labels to 1..K (no gaps after size-filter)', () => { + const H = 64, + W = 64; + // Three valid clusters in different corners. + const clusters = [ + { converge: { y: 12, x: 12 }, source: { y: 4, x: 4 } }, + { converge: { y: 12, x: 52 }, source: { y: 4, x: 44 } }, + { converge: { y: 52, x: 12 }, source: { y: 44, x: 4 } }, + ]; + const { pFinal, seedY, seedX } = makeClusters(H, W, clusters, 25); + const r = getMasks(pFinal, seedY, seedX, H, W); + // Labels must be exactly {1, 2, 3} with no gaps. + const labelSet = new Set(); + for (let i = 0; i < r.masks.length; i++) { + if ((r.masks[i] as number) > 0) labelSet.add(r.masks[i] as number); + } + expect([...labelSet].sort((a, b) => a - b)).toEqual([1, 2, 3]); + }); + + it('handles negative pre-padded coordinates by clamping to 0', () => { + // Edge case: a seed converges past the top-left corner. After +RPAD it + // should still land at row/col 0 of the padded histogram. + const H = 32, + W = 32; + const n = 30; + const pFinal = new Int32Array(2 * n); + const seedY = new Int32Array(n); + const seedX = new Int32Array(n); + for (let i = 0; i < n; i++) { + seedY[i] = i; + seedX[i] = i; + pFinal[2 * i] = -10; // would be -10 + 20 = 10 (still ok) + pFinal[2 * i + 1] = -25; // would be -25 + 20 = -5 → clamped to 0 + } + const r = getMasks(pFinal, seedY, seedX, H, W); + // Just verify no crash and a sensible label produced. + expect(r.count).toBeGreaterThanOrEqual(0); + }); +}); diff --git a/tests/percentile.test.ts b/tests/percentile.test.ts new file mode 100644 index 0000000..2e2dfd2 --- /dev/null +++ b/tests/percentile.test.ts @@ -0,0 +1,62 @@ +/** + * Unit tests for the percentile() helper (preprocess/normalize.ts). + * + * percentile() is the foundation of normalize99 — anchoring it with direct + * tests catches regressions that would otherwise only show up indirectly in + * the normalize99 fixture diffs. + * + * Reference values match numpy.percentile(..., method='linear') which is + * what cellpose's normalize99 calls. + */ +import { describe, it, expect } from 'vitest'; +import { percentile } from '../src/preprocess/normalize.js'; + +describe('percentile (numpy linear interpolation)', () => { + it('returns the only element for a singleton array', () => { + expect(percentile([42], 50)).toBe(42); + expect(percentile([42], 0)).toBe(42); + expect(percentile([42], 99)).toBe(42); + }); + + it('returns NaN for empty input', () => { + expect(Number.isNaN(percentile([], 50))).toBe(true); + }); + + it('p=0 is min, p=100 is max', () => { + const data = [5, 3, 8, 1, 9, 2]; + expect(percentile(data, 0)).toBe(1); + expect(percentile(data, 100)).toBe(9); + }); + + it('median (p=50) matches numpy on a uniform 1..10', () => { + // np.percentile(range(1, 11), 50) = 5.5 + const data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; + expect(percentile(data, 50)).toBeCloseTo(5.5, 6); + }); + + it('p=1 and p=99 on 0..99 match numpy linear interpolation', () => { + // np.arange(100); np.percentile(x, 1) = 0.99; np.percentile(x, 99) = 98.01 + const data = Array.from({ length: 100 }, (_, i) => i); + expect(percentile(data, 1)).toBeCloseTo(0.99, 6); + expect(percentile(data, 99)).toBeCloseTo(98.01, 6); + }); + + it('sorts before computing — unordered input gives same answer', () => { + const a = [1, 2, 3, 4, 5]; + const b = [5, 1, 3, 4, 2]; + expect(percentile(a, 25)).toBe(percentile(b, 25)); + expect(percentile(a, 75)).toBe(percentile(b, 75)); + }); + + it('does not mutate the input (works on a copy)', () => { + const data = [5, 3, 8, 1, 9, 2]; + const snap = [...data]; + percentile(data, 50); + expect(data).toEqual(snap); + }); + + it('accepts Float32Array input', () => { + const f = Float32Array.from([0, 0.25, 0.5, 0.75, 1.0]); + expect(percentile(f, 50)).toBeCloseTo(0.5, 6); + }); +}); diff --git a/tests/resize.test.ts b/tests/resize.test.ts new file mode 100644 index 0000000..13a8a50 --- /dev/null +++ b/tests/resize.test.ts @@ -0,0 +1,126 @@ +/** + * Direct correctness tests for diameterResize's bilinear interpolation. + * + * The pure-JS bilinear (resize.ts) replaced the earlier canvas-based path + * that quantized to 8-bit. These tests anchor the math against: + * - identity (no resize), + * - 2× downscale of a known pattern (each dst pixel is the mean of a 2×2 + * src block under the pixel-center mapping), + * - 2× upscale of a known pattern (corners replicate edge, middles + * interpolate), + * - per-channel independence (channels don't bleed into each other), + * - scale return value matches Python's image_scaling = 30 / diameter. + */ +import { describe, it, expect } from 'vitest'; +import { diameterResize } from '../src/preprocess/resize.js'; + +describe('diameterResize bilinear (pure JS)', () => { + it('identity: srcSize == dstSize → values bit-exact match input', () => { + const W = 8, + H = 8; + const chw = new Float32Array(3 * H * W); + for (let i = 0; i < chw.length; i++) chw[i] = i; + // diameter = targetDiameter = 30 → fast-path scale=1 copy. + const r = diameterResize(chw, W, H, { channels: 3, diameter: 30 }); + expect(r.scale).toBe(1); + expect(r.width).toBe(W); + expect(r.height).toBe(H); + for (let i = 0; i < chw.length; i++) expect(r.data[i]).toBe(chw[i]); + }); + + it('2× downscale: each dst pixel = mean of corresponding 2×2 src block', () => { + // 4×4 src with values 0..15 row-major (single channel). + const srcW = 4, + srcH = 4; + const chw = new Float32Array(srcW * srcH); + for (let i = 0; i < chw.length; i++) chw[i] = i; + // diameter=60, target=30 → scale = 0.5 → dst 2x2. + const r = diameterResize(chw, srcW, srcH, { channels: 1, diameter: 60 }); + expect(r.scale).toBe(0.5); + expect(r.width).toBe(2); + expect(r.height).toBe(2); + // dst[0,0] samples sy = 0.5, sx = 0.5 → corners (0,0)(0,1)(1,0)(1,1) at 25% + // each → mean(0, 1, 4, 5) = 2.5. + expect(r.data[0]).toBeCloseTo(2.5, 5); + // dst[0,1]: sy=0.5, sx=2.5 → mean(2, 3, 6, 7) = 4.5. + expect(r.data[1]).toBeCloseTo(4.5, 5); + // dst[1,0]: sy=2.5, sx=0.5 → mean(8, 9, 12, 13) = 10.5. + expect(r.data[2]).toBeCloseTo(10.5, 5); + // dst[1,1]: sy=2.5, sx=2.5 → mean(10, 11, 14, 15) = 12.5. + expect(r.data[3]).toBeCloseTo(12.5, 5); + }); + + it('2× upscale: corners replicate, interior interpolates linearly', () => { + // 2×2 src: row 0 = [0, 10], row 1 = [20, 30]. Single channel. + const chw = new Float32Array([0, 10, 20, 30]); + // diameter=15, target=30 → scale=2 → dst 4x4. + const r = diameterResize(chw, 2, 2, { channels: 1, diameter: 15 }); + expect(r.scale).toBe(2); + expect(r.width).toBe(4); + expect(r.height).toBe(4); + // Top-left dst corner samples sy=sx=-0.25; both y0/y1 and x0/x1 clamp to + // 0 (edge replication) → result == src[0,0] == 0. + expect(r.data[0]).toBeCloseTo(0, 5); + // Top-right dst corner: sy=-0.25, sx=1.25 → y clamps to row 0, x to + // col 1 (also clamps since x1=2 is out of bounds) → result == src[0,1] == 10. + expect(r.data[3]).toBeCloseTo(10, 5); + // Bottom-left corner: → src[1,0] == 20. + expect(r.data[12]).toBeCloseTo(20, 5); + // Bottom-right corner: → src[1,1] == 30. + expect(r.data[15]).toBeCloseTo(30, 5); + // Interior dst[1, 1] samples sy=0.25, sx=0.25 → fy=fx=0.25 → weighted + // sum: 0*0.75² + 10*0.75*0.25 + 20*0.25*0.75 + 30*0.25² = 1.875 + 3.75 + // + 1.875 = 7.5. (Also the value of the linear extrapolation 10x + 20y + // at (0.25, 0.25), which is the sanity check: bilinear is exact for + // linear functions.) + expect(r.data[5]).toBeCloseTo(7.5, 5); + }); + + it('channels are independent (no bleed between channel buffers)', () => { + // 2x2 image, 3 channels: each channel filled with a distinct constant. + const chw = new Float32Array(3 * 4); + chw.fill(100, 0, 4); + chw.fill(200, 4, 8); + chw.fill(50, 8, 12); + // 2x upscale: each channel of the dst should still be (almost) constant. + const r = diameterResize(chw, 2, 2, { channels: 3, diameter: 15 }); + expect(r.width).toBe(4); + expect(r.height).toBe(4); + // dst layout is also CHW: channel c occupies [c*hwOut, (c+1)*hwOut). + for (let i = 0; i < 16; i++) expect(r.data[i]).toBeCloseTo(100, 5); + for (let i = 16; i < 32; i++) expect(r.data[i]).toBeCloseTo(200, 5); + for (let i = 32; i < 48; i++) expect(r.data[i]).toBeCloseTo(50, 5); + }); + + it('scale return value equals 30 / diameter (cellpose convention)', () => { + const chw = new Float32Array(3 * 32 * 32); + expect(diameterResize(chw, 32, 32, { channels: 3, diameter: 60 }).scale).toBeCloseTo(0.5, 6); + expect(diameterResize(chw, 32, 32, { channels: 3, diameter: 15 }).scale).toBeCloseTo(2.0, 6); + expect(diameterResize(chw, 32, 32, { channels: 3, diameter: 30 }).scale).toBe(1); + }); + + it('FP precision: constant input round-trips to FP-rounding tolerance', () => { + // Motivation: fix #2 replaced the canvas-based resize that quantized to + // uint8 (~1/255 = 0.0039 error per pixel, AFTER rescaling back to the + // source range). A constant-valued image through that path picked up + // ~0.004 error from the 8-bit quantization. With pure-JS bilinear on + // Float32, a constant input stays constant to FP-rounding noise. + const W = 24, + H = 24; + const src = new Float32Array(W * H); + src.fill(0.42); + const down = diameterResize(src, W, H, { channels: 1, diameter: 60 }); + const up = diameterResize(down.data, down.width, down.height, { + channels: 1, + diameter: 15, + }); + let maxAbs = 0; + for (let i = 0; i < up.data.length; i++) { + maxAbs = Math.max(maxAbs, Math.abs((up.data[i] as number) - 0.42)); + } + // The 1/255 floor is 0.0039. Anything well below that proves no + // 8-bit-quantization step is in the path. FP roundtrip on a constant is + // typically <1e-7. + expect(maxAbs).toBeLessThan(1e-5); + }); +}); diff --git a/tests/resize_guard.test.ts b/tests/resize_guard.test.ts new file mode 100644 index 0000000..e34ffe9 --- /dev/null +++ b/tests/resize_guard.test.ts @@ -0,0 +1,58 @@ +/** + * Unit tests for diameterResize parameter checks (preprocess/resize.ts). + * + * The real bilinear resize uses OffscreenCanvas / HTMLCanvas, which aren't + * available in the Node-based test runner. These tests cover the + * pre-canvas paths only: invalid diameter, the no-op pass-through when + * scale ≈ 1, and the 4096×4096 memory guard added in M7. + */ +import { describe, it, expect } from 'vitest'; +import { diameterResize } from '../src/preprocess/resize.js'; + +describe('diameterResize parameter validation', () => { + it('throws when diameter <= 0', () => { + const img = new Float32Array(3 * 64 * 64); + expect(() => diameterResize(img, 64, 64, { channels: 3, diameter: 0 })).toThrow( + /diameter must be positive/, + ); + expect(() => diameterResize(img, 64, 64, { channels: 3, diameter: -5 })).toThrow( + /diameter must be positive/, + ); + }); + + it('returns a copy + scale=1 when source diameter already matches target', () => { + // diameter = targetDiameter (default 30) → scale = 1.0 → no canvas op. + const img = new Float32Array(3 * 32 * 32); + for (let i = 0; i < img.length; i++) img[i] = i; + const r = diameterResize(img, 32, 32, { channels: 3, diameter: 30 }); + expect(r.scale).toBe(1); + expect(r.width).toBe(32); + expect(r.height).toBe(32); + // Distinct buffer (no aliasing). + expect(r.data).not.toBe(img); + expect(r.data.length).toBe(img.length); + for (let i = 0; i < img.length; i++) expect(r.data[i]).toBe(img[i]); + }); + + it('enforces the 4096x4096 output cap before allocating canvases', () => { + // Source 1000×1000, asked for a 12.5x upscale → 12500×12500 destination. + // That exceeds 4096×4096 — should throw with a clear message. + const img = new Float32Array(3 * 1000 * 1000); + expect(() => diameterResize(img, 1000, 1000, { channels: 3, diameter: 2.4 })).toThrow( + /exceeds the safe limit/, + ); + }); + + it('memory-guard error includes the requested scale and source size', () => { + const img = new Float32Array(3 * 2000 * 2000); + let caught: Error | null = null; + try { + diameterResize(img, 2000, 2000, { channels: 3, diameter: 3 }); + } catch (e) { + caught = e as Error; + } + expect(caught).not.toBeNull(); + expect(caught!.message).toMatch(/2000x2000/); + expect(caught!.message).toMatch(/diameter is 3/); + }); +}); diff --git a/tests/taper_mask.test.ts b/tests/taper_mask.test.ts new file mode 100644 index 0000000..9ef3bd5 --- /dev/null +++ b/tests/taper_mask.test.ts @@ -0,0 +1,99 @@ +/** + * Unit tests for taperMask (postprocess/average_tiles.ts). + * + * The taper mask is a 2D outer-product of a sigmoid window. It should be: + * - Symmetric around the center on each axis. + * - Maximal at the center (where the sigmoid saturates at ~1). + * - Close to zero in the 20-pixel border (sigmoid inflection point). + * - Cached per B so repeated calls return the same array. + * + * Matches cellpose.transforms._taper_mask formulation: + * bsize = max(224, B); + * xm = abs(arange(bsize) - mean) + * m1d = 1 / (1 + exp((xm - (bsize/2 - 20)) / 7.5)) + * mask = outer(m1d, m1d) + * then cropped to (B, B). + */ +import { describe, it, expect } from 'vitest'; +import { taperMask } from '../src/postprocess/average_tiles.js'; + +describe('taperMask', () => { + it('returns the same array on repeat calls (caching)', () => { + const a = taperMask(256); + const b = taperMask(256); + expect(a).toBe(b); + }); + + it('has correct shape (B*B floats)', () => { + expect(taperMask(256).length).toBe(256 * 256); + expect(taperMask(128).length).toBe(128 * 128); + }); + + it('peaks near the center (~1.0)', () => { + const B = 256; + const m = taperMask(B); + const center = m[(B / 2) * B + B / 2] as number; + // Sigmoid saturates near 1 in the center. + expect(center).toBeGreaterThan(0.99); + expect(center).toBeLessThanOrEqual(1.0); + }); + + it('is small in the 20-px border (well under 0.5)', () => { + const B = 256; + const m = taperMask(B); + // The sigmoid inflection is at bsize/2 - 20 ≈ 108 (for B=256, since + // bsize = max(224, 256) = 256). So at the very corner (0, 0) we're + // 128 px from center → 20 px past inflection → sigmoid is ~exp(-20/7.5) + // ≈ 0.07, but the OUTER product squares this, so ~0.005. + const corner = m[0] as number; + expect(corner).toBeLessThan(0.05); + expect(corner).toBeGreaterThan(0); + }); + + it('is symmetric on both axes', () => { + const B = 64; + const m = taperMask(B); + for (let y = 0; y < B; y++) { + for (let x = 0; x < B; x++) { + const a = m[y * B + x] as number; + const b = m[y * B + (B - 1 - x)] as number; + expect(a).toBeCloseTo(b, 6); + } + } + for (let y = 0; y < B; y++) { + for (let x = 0; x < B; x++) { + const a = m[y * B + x] as number; + const b = m[(B - 1 - y) * B + x] as number; + expect(a).toBeCloseTo(b, 6); + } + } + }); + + it('matches Python formula at a handful of points (B=256)', () => { + const B = 256; + const m = taperMask(B); + const bsize = Math.max(224, B); // 256 + const mean = (bsize - 1) / 2; + const inflection = bsize / 2 - 20; + const sigma = 7.5; + // Cropping offset: lo = (bsize - B) >> 1 → 0 when bsize == B. + const lo = (bsize - B) >> 1; + const sigmoid = (i: number): number => { + const x = lo + i; + return 1 / (1 + Math.exp((Math.abs(x - mean) - inflection) / sigma)); + }; + // Sample five (y, x) positions. + const samples: Array<[number, number]> = [ + [0, 0], + [10, 10], + [50, 200], + [128, 128], + [200, 50], + ]; + for (const [y, x] of samples) { + const expected = sigmoid(y) * sigmoid(x); + const got = m[y * B + x] as number; + expect(got).toBeCloseTo(expected, 6); + } + }); +}); From 7ccf0c2fcb87f03cb1452957d9655d04455d3286 Mon Sep 17 00:00:00 2001 From: belkassaby Date: Fri, 22 May 2026 21:48:33 -0400 Subject: [PATCH 2/5] fix: make demo work in vite dev mode; abortable model preload; status hooks Three issues surfaced while smoke-testing the demo: 1. Image picker silently failed. `loadImageFile` uses `createImageBitmap`, which throws on formats it can't decode (TIFF microscopy images, etc). The change handler had no try/catch so the rejection was swallowed and the input canvas stayed blank. Now: try/catch around the load + draw + `log()` the error visibly. The `createImageBitmap` catch path also nudges the user toward PNG/JPEG. 2. Worker module failed to load with no diagnostics. Vite's worker plugin transforms `new URL('./inference.worker.js', import.meta.url)` into a URL ending in `inference.worker.js?worker_file`. The library source uses `.js` so the tsc-built prod artifact resolves correctly, but in dev the file on disk is `.ts` and the worker plugin doesn't auto-swap extensions like the regular resolver does. Vite then returned the SPA fallback (200 OK index.html with `text/html` content-type), which the browser failed to parse as a JS module - silently, because the existing `worker.onerror` handler only rejected pending tile promises (empty during init). Two fixes: (a) add a small middleware to the demo's vite.config.ts that rewrites `.js?worker_file` -> `.ts?worker_file` when the .ts source exists. Dev-only; prod build unaffected. (b) make `worker.onerror`, `worker.messageerror`, and worker-side `error` messages reject the `_workerReady` promise during init, instead of being dropped. 3. Abort didn't cancel a stuck model load. The abort signal was only wired to `segment()`, not to `Cellpose.fromPretrained(...)`. If the user hit Abort during the preload phase (fetch / IDB write / ORT session create) nothing happened. Now: `FromPretrainedOptions.signal` propagates into `_ensureWorker`, where an abort terminates the worker and rejects the init promise with `AbortError`. Bonus: new `FromPretrainedOptions.onStatus` callback. The worker now posts pre-ready status strings ('worker module loaded', 'init message received', 'configuring ORT', 'creating ORT session', 'session created in '), and the demo logs each with timing info. Removes the "fetch:100% then silence" mystery. Verified end-to-end in Firefox: 462x346 image -> 6 tiles, 626 ms/tile median, 75 masks, 4.26 s total. --- examples/demo/main.ts | 53 +++++++++++++++++----- examples/demo/vite.config.ts | 38 +++++++++++++++- src/cellpose.ts | 86 +++++++++++++++++++++++++++++++++--- src/inference.worker.ts | 19 +++++++- src/worker-protocol.ts | 6 ++- 5 files changed, 180 insertions(+), 22 deletions(-) diff --git a/examples/demo/main.ts b/examples/demo/main.ts index e105f7b..6f74708 100644 --- a/examples/demo/main.ts +++ b/examples/demo/main.ts @@ -20,7 +20,7 @@ const fmt = (ms: number) => (ms < 1000 ? `${ms.toFixed(0)} ms` : `${(ms / 1000). let cp: Cellpose | null = null; let currentImage: SegmentInput | null = null; -async function ensureModel(): Promise { +async function ensureModel(signal?: AbortSignal): Promise { if (cp) return cp; const modelUrl = $('modelUrl').value; const preload = ($('preload') as HTMLInputElement).checked; @@ -28,18 +28,26 @@ async function ensureModel(): Promise { log(`fromPretrained('${modelUrl}', { preload: ${preload}, bypassCache: ${bypassCache} })...`); const t0 = performance.now(); let lastPctLogged = -1; - cp = await Cellpose.fromPretrained(modelUrl, { + const fetchT0 = { v: 0 }; + const opts: Parameters[1] = { preload, bypassCache, onProgress: ({ loaded, total }) => { if (!total) return; const pct = Math.floor((loaded / total) * 100); if (pct >= lastPctLogged + 10 || pct === 100) { - log(` fetch: ${pct}%`); + log(` fetch: ${pct}% (${fmt(performance.now() - t0)})`); lastPctLogged = pct; + if (pct === 100) fetchT0.v = performance.now(); } }, - }); + onStatus: (s) => { + const since = fetchT0.v ? ` (+${fmt(performance.now() - fetchT0.v)} since fetch:100%)` : ''; + log(` worker: ${s}${since}`); + }, + }; + if (signal) opts.signal = signal; + cp = await Cellpose.fromPretrained(modelUrl, opts); log(`model ready in ${fmt(performance.now() - t0)}`); return cp; } @@ -156,7 +164,19 @@ function drawHeatmap(canvas: HTMLCanvasElement, data: Float32Array, w: number, h } async function loadImageFile(file: File): Promise { - const bitmap = await createImageBitmap(file); + // `createImageBitmap` is the fast path (JPEG, PNG, WebP, GIF, BMP, ICO, + // SVG in most browsers). Microscopy users often have TIFF — that format + // isn't decoded natively by either createImageBitmap or HTMLImageElement, + // so we surface a clear error in the caller's try/catch. + let bitmap: ImageBitmap; + try { + bitmap = await createImageBitmap(file); + } catch (err) { + throw new Error( + `couldn't decode ${file.type || file.name}: ${(err as Error).message}. ` + + `Browsers don't natively decode TIFF — convert to PNG or JPEG first.`, + ); + } const canvas = new OffscreenCanvas(bitmap.width, bitmap.height); const ctx = canvas.getContext('2d')!; ctx.drawImage(bitmap, 0, 0); @@ -200,7 +220,9 @@ async function run() { ($('abort') as HTMLButtonElement).disabled = false; activeAbort = new AbortController(); try { - const model = await ensureModel(); + // Pass the signal so abort also cancels a stuck model load (fetch, IDB + // write, or ORT session create), not just segment(). + const model = await ensureModel(activeAbort.signal); const adapter = await model.describeAdapter(); log(`adapter: vendor=${adapter?.vendor} arch=${adapter?.architecture}`); @@ -261,12 +283,19 @@ async function run() { $('imageFile').addEventListener('change', async (ev) => { const f = (ev.target as HTMLInputElement).files?.[0]; if (!f) return; - currentImage = await loadImageFile(f); - drawToCanvas($('inputCanvas') as HTMLCanvasElement, currentImage); - ($('go') as HTMLButtonElement).disabled = false; - log( - `loaded ${f.name}: ${currentImage.width}x${currentImage.height}, ch=${currentImage.channels}`, - ); + try { + currentImage = await loadImageFile(f); + drawToCanvas($('inputCanvas') as HTMLCanvasElement, currentImage); + ($('go') as HTMLButtonElement).disabled = false; + log( + `loaded ${f.name} (${f.type || 'unknown type'}): ` + + `${currentImage.width}x${currentImage.height}, ch=${currentImage.channels}`, + ); + } catch (err) { + const e = err as Error; + log(`failed to load ${f.name}: ${e.message}`, 'fail'); + console.error(err); + } }); $('useSynthetic').addEventListener('click', () => { currentImage = makeSynthetic(); diff --git a/examples/demo/vite.config.ts b/examples/demo/vite.config.ts index 53b1ec7..d868b06 100644 --- a/examples/demo/vite.config.ts +++ b/examples/demo/vite.config.ts @@ -1,4 +1,5 @@ -import { defineConfig } from 'vite'; +import { defineConfig, type Plugin } from 'vite'; +import fs from 'node:fs'; /** * The model itself is served from examples/demo/public/cpsam_fp16.onnx @@ -14,7 +15,42 @@ import { defineConfig } from 'vite'; * also blocked. Solution: proxy /ort/* to jsdelivr — the browser sees * same-origin URLs while bytes come from jsdelivr's CORS-enabled CDN. */ + +/** + * Vite's worker plugin transforms `new URL('./inference.worker.js', + * import.meta.url)` into a URL ending in `inference.worker.js?worker_file`. + * The library source uses `.js` so the tsc-built prod artifact resolves + * correctly. But in dev mode the file on disk is `.ts`, and the worker + * plugin doesn't auto-swap extensions like the regular resolver does — the + * request 404s, then falls through to the SPA fallback (200 OK index.html), + * and the browser silently fails to parse HTML as a JS module. + * + * This middleware rewrites `…/foo.js?worker_file…` → `…/foo.ts?worker_file…` + * when the `.ts` file exists on disk. Dev-only; the prod build path is + * unaffected. + */ +const remapWorkerTsExt = (): Plugin => ({ + name: 'cellpose-demo-remap-worker-ts-ext', + configureServer(server) { + server.middlewares.use((req, _res, next) => { + const url = req.url; + if (url && url.includes('worker_file') && url.includes('.js?')) { + const [pathPart, queryPart = ''] = url.split('?'); + if (pathPart && pathPart.endsWith('.js')) { + const tsPath = pathPart.replace(/\.js$/, '.ts'); + const onDisk = tsPath.replace(/^\/@fs/, ''); + if (fs.existsSync(onDisk)) { + req.url = `${tsPath}?${queryPart}`; + } + } + } + next(); + }); + }, +}); + export default defineConfig({ + plugins: [remapWorkerTsExt()], server: { proxy: { '/ort': { diff --git a/src/cellpose.ts b/src/cellpose.ts index de551d2..4a1d1ac 100644 --- a/src/cellpose.ts +++ b/src/cellpose.ts @@ -25,7 +25,13 @@ export interface FromPretrainedOptions { /** Forwarded to the model fetcher. */ onProgress?: (p: FetchProgress) => void; bypassCache?: boolean; + /** Aborts the in-flight fetch AND any preload worker init. */ signal?: AbortSignal; + /** Optional callback for worker-init phase strings ('spawning worker', + * 'creating ORT session', 'session ready'). Useful for showing a determinate + * status while the 588 MB FP16 model is being parsed and the WebGPU adapter + * is being initialized. */ + onStatus?: (status: string) => void; } export interface SegmentTileOutput { @@ -120,9 +126,15 @@ export class Cellpose { ...(opts.signal !== undefined && { signal: opts.signal }), }; const bytes = await fetchModel(modelUrl, fetchOpts); + if (opts.signal?.aborted) throw new DOMException('Aborted', 'AbortError'); if (opts.wasmPaths) configureOrt({ wasmPaths: opts.wasmPaths }); const cp = new Cellpose(bytes, modelUrl); - if (opts.preload) await cp._ensureWorker(); + if (opts.preload) { + const workerOpts: { signal?: AbortSignal; onStatus?: (s: string) => void } = {}; + if (opts.signal !== undefined) workerOpts.signal = opts.signal; + if (opts.onStatus !== undefined) workerOpts.onStatus = opts.onStatus; + await cp._ensureWorker(workerOpts); + } return cp; } @@ -136,13 +148,25 @@ export class Cellpose { } /** Lazy spawn + init. Idempotent. */ - private _ensureWorker(): Promise { + private _ensureWorker( + opts: { signal?: AbortSignal; onStatus?: (s: string) => void } = {}, + ): Promise { if (this._workerReady) return this._workerReady; + if (opts.signal?.aborted) { + return Promise.reject(new DOMException('Aborted before worker init', 'AbortError')); + } + opts.onStatus?.('spawning worker'); const worker = new Worker(new URL('./inference.worker.js', import.meta.url), { type: 'module', }); 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 + // ref is updated to a no-op once init completes. + let rejectInit: (err: Error) => void = () => {}; + let initSettled = false; worker.addEventListener('message', (ev: MessageEvent) => { const msg = ev.data; @@ -159,28 +183,73 @@ export class Cellpose { this._pending.delete(msg.tileId); p.reject(new Error(msg.message)); } + } else if (!initSettled) { + rejectInit(new Error(msg.message)); } + } else if (msg.type === 'status') { + opts.onStatus?.(msg.status); } }); + // The Worker.onerror event fires when the worker module fails to load + // or throws at top level. `ev.message` carries the underlying Error. worker.addEventListener('error', (ev) => { - const err = new Error(ev.message || 'worker error'); + const detail = + ev.message || + (ev.filename ? `worker error in ${ev.filename}:${ev.lineno}` : 'worker error (no detail)'); + 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(); + }); + // 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._workerReady = new Promise((resolve, reject) => { + let abortListener: (() => void) | null = null; + const cleanup = (): void => { + initSettled = true; + if (abortListener && opts.signal) { + opts.signal.removeEventListener('abort', abortListener); + } + worker.removeEventListener('message', onReady); + }; + rejectInit = (err) => { + cleanup(); + reject(err); + }; const onReady = (ev: MessageEvent) => { if (ev.data.type === 'ready') { this._adapterInfo = ev.data.adapterInfo; - worker.removeEventListener('message', onReady); + cleanup(); resolve(); } else if (ev.data.type === 'error' && ev.data.tileId === null) { - worker.removeEventListener('message', onReady); + cleanup(); reject(new Error(ev.data.message)); } }; worker.addEventListener('message', onReady); + // Abort handling: terminate the worker and reject. Honors caller abort + // while ORT is doing its (potentially slow) WebGPU + WASM-sidecar init. + if (opts.signal) { + abortListener = () => { + cleanup(); + this._worker?.terminate(); + this._worker = null; + this._workerReady = null; + reject(new DOMException('Aborted during worker init', 'AbortError')); + }; + opts.signal.addEventListener('abort', abortListener); + } + // Resolve model bytes. If the previous worker took ownership (transfer // detached the buffer), refetch from IDB cache — typically <100 ms. const ensureBytes = async (): Promise => { @@ -191,7 +260,9 @@ export class Cellpose { }; ensureBytes() .then((bytes) => { + if (opts.signal?.aborted) return; // abort fired during fetch; listener already rejected this._modelBytes = null; // drop our ref since we're about to transfer ownership + opts.onStatus?.('posting model to worker'); const init: MainToWorker = { type: 'init', modelBytes: bytes, @@ -199,7 +270,10 @@ export class Cellpose { }; worker.postMessage(init, [bytes]); }) - .catch(reject); + .catch((err) => { + cleanup(); + reject(err); + }); }); return this._workerReady; } diff --git a/src/inference.worker.ts b/src/inference.worker.ts index 6e4fb08..4061ea0 100644 --- a/src/inference.worker.ts +++ b/src/inference.worker.ts @@ -43,11 +43,18 @@ async function describeAdapter(): Promise<{ } async function handleInit(msg: Extract): Promise { + postReply({ type: 'status', status: 'configuring ORT (loading WASM sidecars)' }); configureOrt(msg.wasmPaths); + postReply({ type: 'status', status: 'creating ORT session (parsing 588 MB ONNX graph)' }); + const t0 = performance.now(); session = await ort.InferenceSession.create(msg.modelBytes, { executionProviders: ['webgpu'], graphOptimizationLevel: 'all', }); + postReply({ + type: 'status', + status: `session created in ${(performance.now() - t0).toFixed(0)} ms; describing adapter`, + }); const adapterInfo = await describeAdapter(); postReply({ type: 'ready', adapterInfo }); } @@ -97,11 +104,19 @@ function postReply(msg: WorkerToMain, transfer?: Transferable[]): void { else self.postMessage(msg); } +// 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). +postReply({ type: 'status', status: 'worker module loaded' }); + self.addEventListener('message', async (ev: MessageEvent) => { const msg = ev.data; try { - if (msg.type === 'init') await handleInit(msg); - else if (msg.type === 'run-tile') await handleRunTile(msg); + if (msg.type === 'init') { + postReply({ type: 'status', status: 'init message received' }); + await handleInit(msg); + } else if (msg.type === 'run-tile') await handleRunTile(msg); else if (msg.type === 'dispose') { session = null; self.close(); diff --git a/src/worker-protocol.ts b/src/worker-protocol.ts index e4612c8..d441904 100644 --- a/src/worker-protocol.ts +++ b/src/worker-protocol.ts @@ -19,4 +19,8 @@ export type MainToWorker = export type WorkerToMain = | { type: 'ready'; adapterInfo: { vendor: string; architecture: string; device: string } | null } | { type: 'tile-result'; tileId: number; flowsCellprob: Float32Array; inferenceMs: number } - | { type: 'error'; tileId: number | null; message: string }; + | { type: 'error'; tileId: number | null; message: string } + // Pre-ready progress strings: 'configuring ORT', 'creating session', + // 'session created', 'describing adapter'. Optional; consumers wire via + // `FromPretrainedOptions.onStatus`. + | { type: 'status'; status: string }; From 6ff9540428babf206b2460f74af826fb426c3543 Mon Sep 17 00:00:00 2001 From: belkassaby Date: Fri, 22 May 2026 21:51:23 -0400 Subject: [PATCH 3/5] fix linting --- tests/follow_flows.test.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/follow_flows.test.ts b/tests/follow_flows.test.ts index 6cf35a7..77dba51 100644 --- a/tests/follow_flows.test.ts +++ b/tests/follow_flows.test.ts @@ -107,8 +107,8 @@ describe('followFlows', () => { const dy = c.y - y; const dx = c.x - x; const r = Math.hypot(dy, dx) || 1; - dP[i] = (dy / r) / 5; - dP[hw + i] = (dx / r) / 5; + dP[i] = dy / r / 5; + dP[hw + i] = dx / r / 5; cellprob[i] = 1.0; } } From 4e631f928c1453a05536953ceccb4d41bee5e242 Mon Sep 17 00:00:00 2001 From: belkassaby Date: Fri, 22 May 2026 22:00:16 -0400 Subject: [PATCH 4/5] chore: bump to 0.2.0 and add CHANGELOG.md 0.2.0 is a pre-1.0 minor bump because `segment()` outputs can shift vs 0.1.1 (Math.trunc fix, preprocess reorder, niter scaling). New additive API surface: `FromPretrainedOptions.onStatus` and `signal` honored during preload. No type-level breaking changes. Consumers who pinned mask outputs in their snapshots should expect drift toward Python parity. --- CHANGELOG.md | 103 +++++++++++++++++++++++++++++++++++++++++++++++++++ package.json | 2 +- 2 files changed, 104 insertions(+), 1 deletion(-) create mode 100644 CHANGELOG.md diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..e57e3d5 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,103 @@ +# Changelog + +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.2.0] — 2026-05-22 + +Python-parity fixes plus quality-of-life additions surfaced by a code review +against `cellpose/{transforms,dynamics,models,core}.py`. Output of `segment()` +may differ from `0.1.1` — values are closer to the upstream reference. Snapshot +or regression-pinned consumers should review. + +### Fixed + +- **Preprocess order.** Normalize now runs **before** resize, matching + `models.py:_run_net`. Resizing first would change the per-channel percentile + statistics that drive normalization. +- **Float32 bilinear resize.** Replaced the canvas-based `diameterResize`, + which quantized each channel to uint8 (~1/255 per-pixel error) and required + `OffscreenCanvas`/`HTMLCanvas`, with a pure-JS implementation using OpenCV's + `INTER_LINEAR` pixel-center mapping + edge replication. Removes the canvas + dependency entirely and is now testable in Node. +- **`niter` scaling.** Now computed as `floor(200 / scale)` when `diameter` is + set, matching Python's `int(200 / image_scaling)`. Upscaled images get fewer + iterations, downscaled get more. Explicit user-supplied `niter` still wins. +- **`follow_flows` rounding.** Switched the final coordinate cast from + `Math.round` to `Math.trunc` to match PyTorch's `.int()`, which truncates + toward zero (it's a C-style cast). The 1-pixel shift previously split two + small overlap-zone labels off their Python counterparts in the + `three_cells_192` fixture. +- **README quickstart example.** `cellprob_threshold: 0` at the top level was + silently ignored — the actual API nests it under `dynamics: + { cellprobThreshold }`. Also removed a stale "once published in M6 + follow-up" parenthetical referencing the model URL. +- **Worker init errors no longer dropped.** `worker.onerror`, + `worker.messageerror`, and worker-side `error` messages received before + the first tile now reject the `_workerReady` promise; previously they were + only routed to pending-tile promises (empty during init) and disappeared + silently. + +### Added + +- **`FromPretrainedOptions.signal` now cancels the preload phase.** Previously + the abort signal was only honored during fetch and `segment()`; if the user + aborted while ORT was creating its session, nothing happened. Now an abort + during `_ensureWorker()` terminates the worker and rejects with + `AbortError`. +- **`FromPretrainedOptions.onStatus`.** New optional callback that receives + pre-ready phase strings from the inference worker (`'worker module loaded'`, + `'init message received'`, `'configuring ORT (loading WASM sidecars)'`, + `'creating ORT session (parsing 588 MB ONNX graph)'`, `'session created in + ms; describing adapter'`). Useful for showing a determinate status + during the slow first-run worker init. +- **Test coverage**: 14 → 61 tests across 11 files. New direct-coverage suites + for `buildCpsamChannels`, `percentile`, `taperMask`, `followFlows`, + `getMasks`, the new bilinear `diameterResize`, plus an end-to-end + postprocess pipeline test that mocks the inference step. + +### Changed + +- **Dynamics parity gate.** The `three_cells_192` fixture's mean-IoU gate + tightened from 0.55 → 0.99 (was a known divergence, now matches Python at + 1.000 with all 5 per-cell IoUs at 1.000). +- **Demo image picker** surfaces errors visibly instead of silently failing. + `createImageBitmap`'s catch path now nudges users toward PNG/JPEG (TIFF + microscopy images aren't natively decoded by any browser). +- **Demo Vite dev mode** now correctly loads the inference worker. Vite's + worker plugin transforms `new URL('./inference.worker.js', ...)` to a URL + ending in `.js?worker_file`, but the source file on disk is `.ts` and the + worker plugin doesn't auto-swap extensions. Demo's `vite.config.ts` now + ships a small middleware that rewrites `.js?worker_file` → + `.ts?worker_file` when the `.ts` source exists. Dev-only; production builds + unaffected. + +### Internal + +- Added `'status'` message type to the worker → main thread protocol. +- Removed canvas dependency from `src/preprocess/resize.ts`. + +## [0.1.1] — 2026-05-15 + +Patch release. + +### Fixed + +- README and docs: corrected Hugging Face Hub model URL to + `ballon999/cellpose-sam-onnx` and normalized URL casing to lowercase across + `docs/PLAN.md`, `docs/MILESTONE7-RESULTS.md`, and `examples/demo/index.html`. + +## [0.1.0] — 2026-05-15 + +Initial public release. End-to-end browser segmentation pipeline: +preprocessing, IndexedDB-cached model load, WebGPU inference in a Web Worker +with `AbortSignal` cancellation, tile averaging, flow dynamics, full-image +label maps. See [`docs/PLAN.md`](./docs/PLAN.md) for the full milestone trail +and [`docs/STAGE0-RESULTS.md`](./docs/STAGE0-RESULTS.md) for parity / latency +measurements. + +[0.2.0]: https://github.com/belkassaby/Cellpose.js/compare/v0.1.1...v0.2.0 +[0.1.1]: https://github.com/belkassaby/Cellpose.js/compare/v0.1.0...v0.1.1 +[0.1.0]: https://github.com/belkassaby/Cellpose.js/releases/tag/v0.1.0 diff --git a/package.json b/package.json index 1534013..e3b7ad8 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "cellpose-js", - "version": "0.1.1", + "version": "0.2.0", "description": "Browser-side cellular segmentation via Cellpose-SAM, running on WebGPU.", "type": "module", "license": "MIT", From 634bc80184a56597d2e35013a4ee4085365ca099 Mon Sep 17 00:00:00 2001 From: belkassaby Date: Fri, 22 May 2026 22:13:56 -0400 Subject: [PATCH 5/5] fix linting --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e57e3d5..aa1fd01 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,7 +32,7 @@ or regression-pinned consumers should review. `three_cells_192` fixture. - **README quickstart example.** `cellprob_threshold: 0` at the top level was silently ignored — the actual API nests it under `dynamics: - { cellprobThreshold }`. Also removed a stale "once published in M6 +{ cellprobThreshold }`. Also removed a stale "once published in M6 follow-up" parenthetical referencing the model URL. - **Worker init errors no longer dropped.** `worker.onerror`, `worker.messageerror`, and worker-side `error` messages received before @@ -51,7 +51,7 @@ or regression-pinned consumers should review. pre-ready phase strings from the inference worker (`'worker module loaded'`, `'init message received'`, `'configuring ORT (loading WASM sidecars)'`, `'creating ORT session (parsing 588 MB ONNX graph)'`, `'session created in - ms; describing adapter'`). Useful for showing a determinate status + ms; describing adapter'`). Useful for showing a determinate status during the slow first-run worker init. - **Test coverage**: 14 → 61 tests across 11 files. New direct-coverage suites for `buildCpsamChannels`, `percentile`, `taperMask`, `followFlows`,