diff --git a/packages/cli/src/background-removal/inference.test.ts b/packages/cli/src/background-removal/inference.test.ts new file mode 100644 index 000000000..ffc62cc0b --- /dev/null +++ b/packages/cli/src/background-removal/inference.test.ts @@ -0,0 +1,18 @@ +import { describe, expect, it } from "vitest"; +import { MEAN, STD } from "./inference.js"; + +// Regression: the u2net_human_seg model was trained with ImageNet +// normalization. Drifting away from these exact values changes the input +// tensor at every pixel and shifts the predicted alpha mask noticeably +// (Miguel reproduced 8,317 pixel changes with delta up to 78/255 when std +// was set to (1, 1, 1)). Reference: +// https://github.com/danielgatis/rembg/blob/main/rembg/sessions/u2net_human_seg.py#L33 +describe("background-removal/inference — rembg u2net_human_seg parity", () => { + it("MEAN matches U2netHumanSegSession reference", () => { + expect(MEAN).toEqual([0.485, 0.456, 0.406]); + }); + + it("STD matches U2netHumanSegSession reference (ImageNet, not the base u2net's (1,1,1))", () => { + expect(STD).toEqual([0.229, 0.224, 0.225]); + }); +}); diff --git a/packages/cli/src/background-removal/inference.ts b/packages/cli/src/background-removal/inference.ts index 3c47c1d31..ced409846 100644 --- a/packages/cli/src/background-removal/inference.ts +++ b/packages/cli/src/background-removal/inference.ts @@ -11,8 +11,12 @@ import { ensureModel, selectProviders, type Device, type ModelId } from "./manag const INPUT_SIZE = 320; const INPUT_PLANE = INPUT_SIZE * INPUT_SIZE; -const MEAN = [0.485, 0.456, 0.406] as const; -const STD = [1.0, 1.0, 1.0] as const; + +// Must match rembg's U2netHumanSegSession.predict — ImageNet mean/std, NOT the +// (1.0, 1.0, 1.0) std used by the general-purpose u2net session. +// https://github.com/danielgatis/rembg/blob/main/rembg/sessions/u2net_human_seg.py#L33 +export const MEAN = [0.485, 0.456, 0.406] as const; +export const STD = [0.229, 0.224, 0.225] as const; type Sharp = typeof sharpType; interface OrtModule { diff --git a/packages/cli/src/background-removal/pipeline.test.ts b/packages/cli/src/background-removal/pipeline.test.ts index 1294e9074..93a870c55 100644 --- a/packages/cli/src/background-removal/pipeline.test.ts +++ b/packages/cli/src/background-removal/pipeline.test.ts @@ -1,5 +1,7 @@ import { describe, expect, it } from "vitest"; -import { inferOutputFormat, inferInputKind, buildEncoderArgs } from "./pipeline.js"; +import { EventEmitter } from "node:events"; +import type { spawn } from "node:child_process"; +import { inferOutputFormat, inferInputKind, buildEncoderArgs, waitForExit } from "./pipeline.js"; describe("background-removal/pipeline — inferOutputFormat", () => { it("maps .webm → webm", () => { @@ -65,3 +67,42 @@ describe("background-removal/pipeline — buildEncoderArgs", () => { expect(args[rIdx + 1]).toBe("24"); }); }); + +// Regression: a previous version of waitForExit treated `code === null` as +// success. Per Node's child_process docs, that's the signal-killed case — +// reporting it as success means a SIGTERM/SIGKILL'd ffmpeg encoder produces +// a "successful" render with a missing or truncated output file. +describe("background-removal/pipeline — waitForExit signal handling", () => { + function fakeProc(): ReturnType { + return new EventEmitter() as unknown as ReturnType; + } + + it("resolves on a clean exit (code=0, signal=null)", async () => { + const proc = fakeProc(); + const promise = waitForExit(proc, "ffmpeg encoder", () => ""); + proc.emit("exit", 0, null); + await expect(promise).resolves.toBeUndefined(); + }); + + it("rejects when killed by signal (code=null, signal='SIGTERM')", async () => { + const proc = fakeProc(); + const promise = waitForExit(proc, "ffmpeg encoder", () => "tail of stderr"); + proc.emit("exit", null, "SIGTERM"); + await expect(promise).rejects.toThrow(/killed by SIGTERM/); + await expect(promise).rejects.toThrow(/tail of stderr/); + }); + + it("rejects on non-zero exit code", async () => { + const proc = fakeProc(); + const promise = waitForExit(proc, "ffmpeg encoder", () => ""); + proc.emit("exit", 1, null); + await expect(promise).rejects.toThrow(/exited with code 1/); + }); + + it("rejects on SIGKILL", async () => { + const proc = fakeProc(); + const promise = waitForExit(proc, "ffmpeg encoder", () => ""); + proc.emit("exit", null, "SIGKILL"); + await expect(promise).rejects.toThrow(/killed by SIGKILL/); + }); +}); diff --git a/packages/cli/src/background-removal/pipeline.ts b/packages/cli/src/background-removal/pipeline.ts index 0bed3e844..c5a7e774e 100644 --- a/packages/cli/src/background-removal/pipeline.ts +++ b/packages/cli/src/background-removal/pipeline.ts @@ -309,16 +309,25 @@ async function runPipeline( return processed; } -function waitForExit( +export function waitForExit( proc: ReturnType, label: string, getStderr: () => string, ): Promise { return new Promise((resolve, reject) => { proc.on("error", reject); - proc.on("exit", (code) => { - if (code === 0 || code === null) resolve(); - else reject(new Error(`${label} exited with code ${code}: ${getStderr().slice(-400)}`)); + // Per Node docs the exit callback is (code, signal): on a normal exit + // `code` is the numeric exit status and `signal` is null; on a + // signal-killed exit `code` is null and `signal` is the signal name. + // Treating null-code as success would silently report SIGTERM/SIGKILL + // as a successful render. + proc.on("exit", (code, signal) => { + if (code === 0 && !signal) { + resolve(); + return; + } + const cause = signal ? `killed by ${signal}` : `exited with code ${code}`; + reject(new Error(`${label} ${cause}: ${getStderr().slice(-400)}`)); }); }); }