Skip repeated and low-motion frames in background removal; add 180-line processing - proof-of-concept/takes

proof-of-concept / takes

Skip repeated and low-motion frames in background removal; add 180-line processing

Jeremy Magland <jmagland@flatironinstitute.org> committed 20 days ago commit 969e06c0acac parent 6ec86a8 Browse files

4 changed files+167−30

README.mdmodified+2−2View file

@@ -48,9 +48,9 @@ The library shows every video with its thumbnail, duration, and description. Dra

4848

4949 **Cropping.** The *Crop* button on a take opens a frame of its screen or camera track with a rectangle that can be moved, resized, or redrawn by dragging; a slider picks the frame to show. This is useful for keeping only one window of a full-screen recording, or for reframing the camera. The rectangle can be free or locked to an aspect ratio, including *Fill layout*, the shape of the area that track occupies in the current layout, so that the cropped screen is not letterboxed. The crop is fit into the layout exactly as the full frame would be, and *Apply to all takes* copies it to the other takes. Like the timeline, cropping is non-destructive: it is stored as fractions of the frame in `video.json`, so it does not depend on the recorded resolution, and background removal does not need to run again when it changes.

5050

51-**Background removal.** In *Layout & export*, the background behind you can be kept, blurred, replaced with a solid color, or removed entirely (*Cutout*, for picture-in-picture, where only you appear over the screen recording, sitting on the bottom edge of the video). Choosing any of these processes the camera track of each take with Robust Video Matting (Lin et al., 2021), a recurrent segmentation model designed for video, which gives soft edges around hair and little flicker between frames. The model runs on the CPU through onnxruntime; takes are processed one at a time in the background, and the result is stored as `cutout.webm` (VP9 with an alpha channel) so that each take is processed only once. The model (about 15 MB) is downloaded from GitHub to `~/.cache/takes/models/` the first time it is needed. Export waits until every take in the timeline has been processed. A running or queued job can be cancelled; a cancelled take is not processed again until you start it.

51+**Background removal.** In *Layout & export*, the background behind you can be kept, blurred, replaced with a solid color, or removed entirely (*Cutout*, for picture-in-picture, where only you appear over the screen recording, sitting on the bottom edge of the video). Choosing any of these processes the camera track of each take with Robust Video Matting (Lin et al., 2021), a recurrent segmentation model designed for video, which gives soft edges around hair and little flicker between frames. The model runs on the CPU through onnxruntime; takes are processed one at a time in the background, and the result is stored as `cutout.webm` (VP9 with an alpha channel) so that each take is processed only once. The model (about 15 MB) is downloaded from GitHub to `~/.cache/takes/models/` the first time it is needed. Export waits until every take in the timeline has been processed. On the CPU, the model is not run on every frame. A frame identical to the one before it gets the same result; webcams often deliver only 15 frames per second in dim light, and the camera track is decoded at 30, so this alone halves the work for such takes. In addition, when the frames on either side of a frame differ little near your outline, the model skips it (never two frames in a row): its matte is the mean of its neighbours' mattes, and its colour is its own camera image with the neighbours' edge corrections, so your face stays exactly as recorded. Compared with running the model on every frame, on a 2-minute stretch of a take with a fair amount of movement, the model ran on 71% of the frames, and in the worst frame 3.8% of the pixels along the outline differed by more than 0.2 in alpha (0.02% on average). The remote GPU path (below) runs the model on every distinct frame, since there the model is not what takes the time. A running or queued job can be cancelled; a cancelled take is not processed again until you start it.

5252

53-**Processing resolution.** The camera is scaled down before background removal to at most the number of lines chosen under *Processing resolution* (240 to 1080, default 480), and the model's coarse pass shrinks with it, since that pass takes most of the time. A small picture-in-picture needs far fewer lines than the camera records: at 720p output, a camera box of a quarter of the width is 320×180 px. On the development laptop, a 426×240 frame took 76 ms, compared with 250 ms for 1280×720. We have not measured how much the matte quality drops at the lowest settings. Changing the setting restarts a running job; takes already processed at a lower resolution offer a *Redo* button, and takes processed at a higher one are kept.

53+**Processing resolution.** The camera is scaled down before background removal to at most the number of lines chosen under *Processing resolution* (180 to 1080, default 480), and the model's coarse pass shrinks with it, since that pass takes most of the time. A small picture-in-picture needs far fewer lines than the camera records: at 720p output, a camera box of a quarter of the width is 320×180 px. On the development laptop (an i7-1355U), the model took 70.7 ms per frame at 1280×720, 16.6 ms at 426×240, and 11.4 ms at 320×180; at 240 lines, a 13-minute take was processed in 9 minutes. We have not measured how much the matte quality drops at the lowest settings. Changing the setting restarts a running job; takes already processed at a lower resolution offer a *Redo* button, and takes processed at a higher one are kept.

5454

5555 **Background removal on another machine.** On a laptop CPU, background removal is slow, so it can instead run on a machine with an NVIDIA GPU that you reach over ssh. Set *ssh host* on the Settings page to an ssh destination such as `office` or `user@host`; it must accept key-based login without prompting, since Takes runs ssh without a terminal. That machine needs a working NVIDIA driver, Python 3.10 or later, and ffmpeg (with libvpx) on the `PATH`. If these come from environment modules or similar, put the commands that load them in *Setup command* (for example `module load python ffmpeg`); they run in a bash login shell before each job. For each take, Takes sends the camera track over a single ssh connection, first scaling it down to the processing resolution (see below) and re-encoding it as H.264 at CRF 18 when that is smaller than the camera; at 240 lines this cut the upload of a 720p camera track by a factor of about 15. It then runs the same model with onnxruntime-gpu in a temporary directory there, receives `cutout.webm`, and removes the temporary directory. The first job creates a Python environment in `~/.cache/takes/venv` on the remote machine (about 2.5 GB, because pip installs the CUDA and cuDNN libraries); delete that directory to rebuild it. For a 60 s, 640×480 take, processing took 27.7 s on a remote RTX 6000 Ada, including the transfer, compared with 536 s on the CPU of a 12-thread laptop; the two cutouts differ by less than 1/255 in alpha on average. (That comparison was made at full resolution, without the re-encoding step, which adds a small difference of its own.)

5656

server/matte-worker.pymodified+12−1View file

@@ -2,7 +2,10 @@

22 # the remote machine together with the camera track and runs it in a temporary directory that holds

33 # in.webm. It runs the same Robust Video Matting model as the local path, through onnxruntime with

44 # the CUDA execution provider, and writes out.webm with the ffmpeg arguments given in the spec, so

5-# the result matches a local run. Stage changes and progress are reported on stderr as

5+# the result matches a local run, except that the local path also skips the model on some frames

6+# with little motion (see matte.ts); here the model is fast, so every distinct frame is processed.

7+# A frame identical to the previous one (as when the decoder repeats frames of a camera that

8+# delivered fewer than 30 per second) gets the same output without running the model. Stage changes and progress are reported on stderr as

69 # "takes:stage <label>" and "takes:progress <frames>" lines.

710

811 import json

@@ -56,10 +59,18 @@ def main():

5659 rgba = np.empty((H, W, 4), np.uint8)

5760 size = W * H * 3

5861 n = 0

62+ prev = None

5963 while True:

6064 buf = decoder.stdout.read(size)

6165 if len(buf) < size:

6266 break

67+ if buf == prev:

68+ encoder.stdin.write(rgba.tobytes())

69+ n += 1

70+ if n % 10 == 0:

71+ report(f'progress {n}')

72+ continue

73+ prev = buf

6374 frame = np.frombuffer(buf, np.uint8).reshape(H, W, 3)

6475 src = (frame.transpose(2, 0, 1)[None] / np.float32(255)).astype(np.float32)

6576 fgr, pha, *rec = sess.run(

server/matte.tsmodified+152−26View file

@@ -4,7 +4,8 @@

44 // cutout.webm (VP9 with an alpha channel), which Chrome plays with transparency and ffmpeg can

55 // composite at export. Takes are processed one at a time, since each job already uses most cores.

66 // The camera is first scaled down to the processing resolution chosen for the video, and a queued

7-// or running job can be cancelled.

7+// or running job can be cancelled. On the CPU, the model is not run on repeated frames or on some

8+// frames with little motion (see SKIP_MOTION).

89

910 import { spawn } from 'node:child_process';

1011 import { once } from 'node:events';

@@ -239,6 +240,66 @@ async function processTake(job: Job, remote: { host: string; setup: string }) {

239240 }

240241 }

241242

243+// Skipping frames on the CPU. Much of a camera track repeats: a webcam that delivers 15 frames per

244+// second in dim light has each frame decoded twice at 30 fps, and when the speaker is still, one

245+// frame differs little from the next. So a frame identical to the previous one gets the same output

246+// without running the model, and when the frames on either side of a frame differ by less than

247+// SKIP_MOTION (mean absolute difference, 0-255 per channel) near the speaker's outline, the model is

248+// not run on it either. Its matte is then the mean of its neighbours' mattes, and its colour is its

249+// own camera image plus the mean of the neighbours' corrections (the model's foreground colour minus

250+// the camera image), which are nonzero only near the outline; so the face stays exactly as recorded.

251+// At most every other frame is skipped this way. On a 13-minute take at 240 lines, with the

252+// threshold of 2, the matte of a skipped frame had at most 3% of its outline pixels off by more than

253+// 0.2 from what the model gives for that frame.

254+const SKIP_MOTION = 2;

255+

256+interface Matted {

257+ raw: Buffer; // RGB, interleaved

258+ pha: Float32Array;

259+ fgr: Float32Array; // planar RGB

260+ band?: Uint8Array; // pixels near the outline, computed when needed

261+}

262+

263+// Pixels within 3 px of where the matte is neither 0 nor 1 (a dilation, done as sliding-window

264+// counts along the rows and then the columns).

265+function outlineBand(pha: Float32Array, W: number, H: number, R = 3) {

266+ const rows = new Uint8Array(W * H);

267+ for (let y = 0; y < H; y++) {

268+ const o = y * W;

269+ const edge = (x: number) => (x >= 0 && x < W && pha[o + x] > 0.02 && pha[o + x] < 0.98 ? 1 : 0);

270+ let count = 0;

271+ for (let x = 0; x < R; x++) count += edge(x);

272+ for (let x = 0; x < W; x++) {

273+ count += edge(x + R) - edge(x - R - 1);

274+ rows[o + x] = count > 0 ? 1 : 0;

275+ }

276+ }

277+ const band = new Uint8Array(W * H);

278+ for (let x = 0; x < W; x++) {

279+ const near = (y: number) => (y >= 0 && y < H ? rows[y * W + x] : 0);

280+ let count = 0;

281+ for (let y = 0; y < R; y++) count += near(y);

282+ for (let y = 0; y < H; y++) {

283+ count += near(y + R) - near(y - R - 1);

284+ band[y * W + x] = count > 0 ? 1 : 0;

285+ }

286+ }

287+ return band;

288+}

289+

290+// Mean absolute difference of two RGB frames over a band (0 when the band is empty).

291+function bandMotion(a: Buffer, b: Buffer, band: Uint8Array) {

292+ let sum = 0;

293+ let n = 0;

294+ for (let i = 0; i < band.length; i++) {

295+ if (!band[i]) continue;

296+ const j = i * 3;

297+ sum += Math.abs(a[j] - b[j]) + Math.abs(a[j + 1] - b[j + 1]) + Math.abs(a[j + 2] - b[j + 2]);

298+ n += 3;

299+ }

300+ return n ? sum / n : 0;

301+}

302+

242303 async function runLocal(job: Job, { src, out, W, H, ds, total }: Plan): Promise<number> {

243304 job.stage = (await exists(MODEL_PATH)) ? 'Loading model' : 'Downloading model';

244305 const sess = await getSession();

@@ -266,41 +327,106 @@ async function runLocal(job: Job, { src, out, W, H, ds, total }: Plan): Promise<

266327 const input = new Float32Array(3 * plane);

267328 let rec: ort.Tensor[] = [0, 1, 2, 3].map(() => new ort.Tensor('float32', new Float32Array([0]), [1, 1, 1, 1]));

268329 const dsr = new ort.Tensor('float32', new Float32Array([ds]), [1]);

269- let n = 0;

330+

331+ const matte = async (raw: Buffer): Promise<Matted> => {

332+ for (let i = 0, j = 0; i < plane; i++, j += 3) {

333+ input[i] = raw[j] / 255;

334+ input[i + plane] = raw[j + 1] / 255;

335+ input[i + 2 * plane] = raw[j + 2] / 255;

336+ }

337+ const r = await sess.run({

338+ src: new ort.Tensor('float32', input, [1, 3, H, W]),

339+ r1i: rec[0],

340+ r2i: rec[1],

341+ r3i: rec[2],

342+ r4i: rec[3],

343+ downsample_ratio: dsr,

344+ });

345+ rec = [r.r1o, r.r2o, r.r3o, r.r4o];

346+ return { raw, pha: r.pha.data as Float32Array, fgr: r.fgr.data as Float32Array };

347+ };

348+ const rgbaOf = (m: Matted) => {

349+ const rgba = new Uint8ClampedArray(plane * 4);

350+ for (let i = 0, j = 0; i < plane; i++, j += 4) {

351+ rgba[j] = m.fgr[i] * 255;

352+ rgba[j + 1] = m.fgr[i + plane] * 255;

353+ rgba[j + 2] = m.fgr[i + 2 * plane] * 255;

354+ rgba[j + 3] = m.pha[i] * 255;

355+ }

356+ return Buffer.from(rgba.buffer);

357+ };

358+ // A skipped frame between two matted ones (see SKIP_MOTION).

359+ const rgbaBetween = (raw: Buffer, p: Matted, q: Matted) => {

360+ const rgba = new Uint8ClampedArray(plane * 4);

361+ for (let i = 0, j = 0, k = 0; i < plane; i++, j += 4, k += 3) {

362+ for (let c = 0; c < 3; c++) {

363+ const corr = p.fgr[i + c * plane] - p.raw[k + c] / 255 + q.fgr[i + c * plane] - q.raw[k + c] / 255;

364+ rgba[j + c] = raw[k + c] + corr * 127.5;

365+ }

366+ rgba[j + 3] = (p.pha[i] + q.pha[i]) * 127.5;

367+ }

368+ return Buffer.from(rgba.buffer);

369+ };

370+

371+ let n = 0; // frames written

372+ const write = async (buf: Buffer, count = 1) => {

373+ for (let c = 0; c < count; c++) {

374+ if (!encoder.stdin.write(buf)) await once(encoder.stdin, 'drain', { signal });

375+ n++;

376+ }

377+ job.progress = Math.min(0.99, n / total);

378+ };

379+

380+ let prevRaw: Buffer | null = null; // the previous decoded frame

381+ let last: Matted | null = null; // the last frame the model ran on

382+ let lastOut: Buffer | null = null; // the output of the previous decoded frame

383+ let pending: { raw: Buffer; count: number } | null = null; // a frame that may be skipped, and its repeats

384+ let processed = 0;

385+ let skipped = 0;

270386 try {

271387 for await (const frame of readFrames(decoder.stdout, plane * 3)) {

272388 signal.throwIfAborted();

273- for (let i = 0, j = 0; i < plane; i++, j += 3) {

274- input[i] = frame[j] / 255;

275- input[i + plane] = frame[j + 1] / 255;

276- input[i + 2 * plane] = frame[j + 2] / 255;

389+ if (prevRaw && frame.equals(prevRaw)) {

390+ if (pending) pending.count++;

391+ else await write(lastOut!);

392+ continue;

277393 }

278- const r = await sess.run({

279- src: new ort.Tensor('float32', input, [1, 3, H, W]),

280- r1i: rec[0],

281- r2i: rec[1],

282- r3i: rec[2],

283- r4i: rec[3],

284- downsample_ratio: dsr,

285- });

286- rec = [r.r1o, r.r2o, r.r3o, r.r4o];

287- const fgr = r.fgr.data as Float32Array;

288- const pha = r.pha.data as Float32Array;

289- const rgba = new Uint8ClampedArray(plane * 4);

290- for (let i = 0, j = 0; i < plane; i++, j += 4) {

291- rgba[j] = fgr[i] * 255;

292- rgba[j + 1] = fgr[i + plane] * 255;

293- rgba[j + 2] = fgr[i + 2 * plane] * 255;

294- rgba[j + 3] = pha[i] * 255;

394+ const raw = Buffer.from(frame); // readFrames reuses its buffer

395+ prevRaw = raw;

396+ if (!last) {

397+ last = await matte(raw);

398+ processed++;

399+ await write((lastOut = rgbaOf(last)));

400+ } else if (!pending) {

401+ pending = { raw, count: 1 };

402+ } else {

403+ last.band ??= outlineBand(last.pha, W, H);

404+ if (bandMotion(last.raw, raw, last.band) < SKIP_MOTION) {

405+ const next = await matte(raw);

406+ processed++;

407+ skipped++;

408+ await write(rgbaBetween(pending.raw, last, next), pending.count);

409+ last = next;

410+ pending = null;

411+ await write((lastOut = rgbaOf(last)));

412+ } else {

413+ last = await matte(pending.raw);

414+ processed++;

415+ await write((lastOut = rgbaOf(last)), pending.count);

416+ pending = { raw, count: 1 };

417+ }

295418 }

296- if (!encoder.stdin.write(Buffer.from(rgba.buffer))) await once(encoder.stdin, 'drain', { signal });

297- n++;

298- job.progress = Math.min(0.99, n / total);

299419 }

300420 signal.throwIfAborted();

421+ if (pending) {

422+ last = await matte(pending.raw);

423+ processed++;

424+ await write(rgbaOf(last), pending.count);

425+ }

301426 encoder.stdin.end();

302427 const [[dcode], [ecode]] = await Promise.all([decoderDone, encoderDone]);

303428 if (dcode !== 0 || ecode !== 0) throw new Error(`ffmpeg failed: ${stderr.trim().split('\n').slice(-4).join('\n')}`);

429+ console.log(`matte: model ran on ${processed} of ${n} frames (${skipped} skipped for low motion, the rest repeated)`);

304430 return n;

305431 } catch (e) {

306432 kill();

shared/types.tsmodified+1−1View file

@@ -48,7 +48,7 @@ export interface Background {

4848 resolution: MatteResolution;

4949 }

5050

51-export const MATTE_RESOLUTIONS = [240, 360, 480, 720, 1080] as const;

51+export const MATTE_RESOLUTIONS = [180, 240, 360, 480, 720, 1080] as const;

5252 export type MatteResolution = (typeof MATTE_RESOLUTIONS)[number];

5353

5454 export const DEFAULT_BACKGROUND: Background = { mode: 'none', color: '#1e2530', blur: 18, resolution: 480 };