From a707b6a882eb22a6c4f7ccc7e8a24f9de582f940 Mon Sep 17 00:00:00 2001 From: James Date: Tue, 5 May 2026 18:18:23 -0700 Subject: [PATCH] fix(cli): pin inverse-alpha invariants, harden encoder stdin MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Extract applyMask helper from postprocess and add 5 unit tests pinning the contract this PR is selling: fg.alpha + bg.alpha === 255 per pixel, RGB triples byte-identical between fg and bg, and bg=null path leaves the bg buffer untouched. Without these, a future postprocess change (mask threshold, premultiplied alpha, gamma) could silently break the inverse-alpha relationship and the existing plumbing tests would all still pass. - Add stdin 'error' listener inside spawnFfmpeg. If either encoder dies mid-render, Node emits an unhandled error on the dead writable on the next .write() and crashes the CLI before waitForExit's reject path can surface the encoder's stderr tail. Doubled encoder count = doubled failure surface, so this is worth pinning down. - Tighten stdio param to a 3-tuple so an accidental 1-element array fails at type-check. - Sharpen backpressure comment: write→true means "highWaterMark not exceeded," not "libuv flushed." Reuse-without-corruption is safe only because session.process is slow enough that libuv drains in between. Addresses review on PR #637. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/background-removal/inference.test.ts | 111 +++++++++++++++++- .../cli/src/background-removal/inference.ts | 54 ++++++--- .../cli/src/background-removal/pipeline.ts | 18 ++- 3 files changed, 164 insertions(+), 19 deletions(-) diff --git a/packages/cli/src/background-removal/inference.test.ts b/packages/cli/src/background-removal/inference.test.ts index ffc62cc0b..719bcab4d 100644 --- a/packages/cli/src/background-removal/inference.test.ts +++ b/packages/cli/src/background-removal/inference.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { MEAN, STD } from "./inference.js"; +import { MEAN, STD, applyMask } from "./inference.js"; // Regression: the u2net_human_seg model was trained with ImageNet // normalization. Drifting away from these exact values changes the input @@ -16,3 +16,112 @@ describe("background-removal/inference — rembg u2net_human_seg parity", () => expect(STD).toEqual([0.229, 0.224, 0.225]); }); }); + +// These tests pin the contract that `--background-output` is built on: +// fg.alpha + bg.alpha === 255 per pixel, and the RGB plane is byte-identical +// between fg and bg. A future change to the postprocess loop (different mask +// threshold, premultiplied alpha, gamma-corrected compositing) that breaks +// either invariant should fail here loudly. +describe("background-removal/inference — applyMask invariants", () => { + function makeRgb(pixels: number): Buffer { + // Deterministic but non-trivial RGB so byte equality is meaningful. + const buf = Buffer.allocUnsafe(pixels * 3); + for (let i = 0; i < pixels; i++) { + buf[i * 3] = (i * 7) & 0xff; + buf[i * 3 + 1] = (i * 13 + 31) & 0xff; + buf[i * 3 + 2] = (i * 19 + 61) & 0xff; + } + return buf; + } + + function makeMask(pixels: number): Buffer { + // Hit the saturation endpoints (0, 255) and a few mid-tone values so the + // 255-m inversion is exercised across the full byte range. + const buf = Buffer.allocUnsafe(pixels); + for (let i = 0; i < pixels; i++) buf[i] = (i * 37) & 0xff; + return buf; + } + + it("dual-output: fg.alpha + bg.alpha === 255 for every pixel", () => { + const pixels = 64; + const rgb = makeRgb(pixels); + const mask = makeMask(pixels); + const fg = Buffer.allocUnsafe(pixels * 4); + const bg = Buffer.allocUnsafe(pixels * 4); + + const result = applyMask(rgb, mask, fg, bg, pixels); + + expect(result.fg).toBe(fg); + expect(result.bg).toBe(bg); + for (let i = 0; i < pixels; i++) { + const sum = fg[i * 4 + 3]! + bg[i * 4 + 3]!; + expect(sum).toBe(255); + } + }); + + it("dual-output: RGB triples are byte-identical between fg and bg", () => { + const pixels = 64; + const rgb = makeRgb(pixels); + const mask = makeMask(pixels); + const fg = Buffer.allocUnsafe(pixels * 4); + const bg = Buffer.allocUnsafe(pixels * 4); + + applyMask(rgb, mask, fg, bg, pixels); + + for (let i = 0; i < pixels; i++) { + expect(fg[i * 4]).toBe(bg[i * 4]); + expect(fg[i * 4 + 1]).toBe(bg[i * 4 + 1]); + expect(fg[i * 4 + 2]).toBe(bg[i * 4 + 2]); + // And both match the source. + expect(fg[i * 4]).toBe(rgb[i * 3]); + expect(fg[i * 4 + 1]).toBe(rgb[i * 3 + 1]); + expect(fg[i * 4 + 2]).toBe(rgb[i * 3 + 2]); + } + }); + + it("dual-output: fg.alpha equals the input mask", () => { + const pixels = 32; + const rgb = makeRgb(pixels); + const mask = makeMask(pixels); + const fg = Buffer.allocUnsafe(pixels * 4); + const bg = Buffer.allocUnsafe(pixels * 4); + + applyMask(rgb, mask, fg, bg, pixels); + + for (let i = 0; i < pixels; i++) { + expect(fg[i * 4 + 3]).toBe(mask[i]); + } + }); + + it("single-output: bg=null returns bg=null and writes only fg", () => { + const pixels = 32; + const rgb = makeRgb(pixels); + const mask = makeMask(pixels); + const fg = Buffer.allocUnsafe(pixels * 4); + + const result = applyMask(rgb, mask, fg, null, pixels); + + expect(result.bg).toBeNull(); + expect(result.fg).toBe(fg); + for (let i = 0; i < pixels; i++) { + expect(fg[i * 4]).toBe(rgb[i * 3]); + expect(fg[i * 4 + 3]).toBe(mask[i]); + } + }); + + it("saturates correctly at mask=0 and mask=255", () => { + // mask=0 → fg.alpha=0 (transparent subject), bg.alpha=255 (fully opaque plate) + // mask=255 → fg.alpha=255 (fully opaque subject), bg.alpha=0 (transparent plate) + const rgb = Buffer.from([10, 20, 30, 40, 50, 60]); + const mask = Buffer.from([0, 255]); + const fg = Buffer.allocUnsafe(8); + const bg = Buffer.allocUnsafe(8); + + applyMask(rgb, mask, fg, bg, 2); + + expect(fg[3]).toBe(0); + expect(bg[3]).toBe(255); + expect(fg[7]).toBe(255); + expect(bg[7]).toBe(0); + }); +}); diff --git a/packages/cli/src/background-removal/inference.ts b/packages/cli/src/background-removal/inference.ts index 614038663..605257fc3 100644 --- a/packages/cli/src/background-removal/inference.ts +++ b/packages/cli/src/background-removal/inference.ts @@ -200,30 +200,50 @@ async function postprocess( .raw() .toBuffer(); - const pixels = width * height; - if (rgbaBgBuf) { + return applyMask(rgb, fullMask, rgbaBuf, rgbaBgBuf, width * height); +} + +/** + * Composite the RGB source frame with the segmentation mask into one or two + * RGBA buffers. The contract this PR is built on: + * - `fg`'s alpha is the mask, `bg`'s alpha (when provided) is `255 − mask`, + * so `fg.alpha + bg.alpha === 255` for every pixel. + * - RGB triples are byte-identical between `fg` and `bg`. + * - When `bg` is null, only `fg` is touched. + * + * Exported for direct unit testing of the invariants above without spinning + * up an ONNX session. + */ +export function applyMask( + rgb: Buffer, + mask: Buffer, + fg: Buffer, + bg: Buffer | null, + pixels: number, +): SessionResult { + if (bg) { for (let i = 0; i < pixels; i++) { const r = rgb[i * 3]!; const g = rgb[i * 3 + 1]!; const b = rgb[i * 3 + 2]!; - const m = fullMask[i]!; + const m = mask[i]!; const o = i * 4; - rgbaBuf[o] = r; - rgbaBuf[o + 1] = g; - rgbaBuf[o + 2] = b; - rgbaBuf[o + 3] = m; - rgbaBgBuf[o] = r; - rgbaBgBuf[o + 1] = g; - rgbaBgBuf[o + 2] = b; - rgbaBgBuf[o + 3] = 255 - m; + fg[o] = r; + fg[o + 1] = g; + fg[o + 2] = b; + fg[o + 3] = m; + bg[o] = r; + bg[o + 1] = g; + bg[o + 2] = b; + bg[o + 3] = 255 - m; } - return { fg: rgbaBuf, bg: rgbaBgBuf }; + return { fg, bg }; } for (let i = 0; i < pixels; i++) { - rgbaBuf[i * 4] = rgb[i * 3]!; - rgbaBuf[i * 4 + 1] = rgb[i * 3 + 1]!; - rgbaBuf[i * 4 + 2] = rgb[i * 3 + 2]!; - rgbaBuf[i * 4 + 3] = fullMask[i]!; + fg[i * 4] = rgb[i * 3]!; + fg[i * 4 + 1] = rgb[i * 3 + 1]!; + fg[i * 4 + 2] = rgb[i * 3 + 2]!; + fg[i * 4 + 3] = mask[i]!; } - return { fg: rgbaBuf, bg: null }; + return { fg, bg: null }; } diff --git a/packages/cli/src/background-removal/pipeline.ts b/packages/cli/src/background-removal/pipeline.ts index cbd60a798..a0a0f3bfc 100644 --- a/packages/cli/src/background-removal/pipeline.ts +++ b/packages/cli/src/background-removal/pipeline.ts @@ -318,12 +318,21 @@ interface FfmpegProc { getStderr: () => string; } -function spawnFfmpeg(args: string[], label: string, stdio: ("ignore" | "pipe")[]): FfmpegProc { +type StdioFd = "ignore" | "pipe"; +type StdioTuple = [StdioFd, StdioFd, StdioFd]; + +function spawnFfmpeg(args: string[], label: string, stdio: StdioTuple): FfmpegProc { const proc = spawn("ffmpeg", args, { stdio }); let stderrBuf = ""; proc.stderr?.on("data", (d: Buffer) => { stderrBuf += d.toString(); }); + // If the encoder dies mid-render, the next .write() to its stdin emits an + // 'error' event on the writable. Without a listener, Node treats it as + // unhandled and crashes the CLI before waitForExit's reject path can + // surface the real cause (encoder stderr tail). Swallowing here is safe — + // the process exit is the source of truth. + proc.stdin?.on("error", () => {}); const exit = waitForExit(proc, label, () => stderrBuf); return { proc, exit, getStderr: () => stderrBuf }; } @@ -383,6 +392,13 @@ async function runPipeline( // Issue both writes before any await so a slow encoder doesn't block // the other. Drain anything that returned false before the next // session.process() — its output buffers are reused per frame. + // + // Subtlety: write() returning true means "highWaterMark not exceeded," + // NOT "libuv has flushed the chunk." The buffer reference is held by + // libuv until the underlying syscall completes. Reusing the session's + // output buffer is safe because the next session.process() call takes + // ~10–50ms (ORT inference) — plenty of event-loop turns for libuv to + // drain. If that ever stops being true, we'd need to copy here. const fgWroteFully = fg.proc.stdin!.write(result.fg); const bgWroteFully = bg && result.bg ? bg.proc.stdin!.write(result.bg) : true; if (!fgWroteFully || !bgWroteFully) {