f3253d0Takes: record, organize, and edit explainer videosJeremy Magland 1// Background removal for camera tracks with Robust Video Matting (Lin et al., 2021), run on the CPU
307b4e2Run background removal on a remote GPU over sshJeremy Magland 2// through onnxruntime, or on a remote machine's GPU over ssh when a host is set in the settings. For each take, camera.webm is decoded to 30 fps RGB frames, the recurrent
f3253d0Takes: record, organize, and edit explainer videosJeremy Magland 3// model predicts a foreground color and an alpha matte per frame, and the result is written as
4// cutout.webm (VP9 with an alpha channel), which Chrome plays with transparency and ffmpeg can
5// composite at export. Takes are processed one at a time, since each job already uses most cores.
be26c78Update matte pipeline, server API and video pageJeremy Magland 6// The camera is first scaled down to the processing resolution chosen for the video, and a queued
969e06cSkip repeated and low-motion frames in background removal; add 180-line processingJeremy Magland 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).
f3253d0Takes: record, organize, and edit explainer videosJeremy Magland 9
10import { spawn } from 'node:child_process';
11import { once } from 'node:events';
307b4e2Run background removal on a remote GPU over sshJeremy Magland 12import { createReadStream, createWriteStream } from 'node:fs';
f3253d0Takes: record, organize, and edit explainer videosJeremy Magland 13import fs from 'node:fs/promises';
14import os from 'node:os';
15import path from 'node:path';
be26c78Update matte pipeline, server API and video pageJeremy Magland 16import { Transform, type Readable } from 'node:stream';
307b4e2Run background removal on a remote GPU over sshJeremy Magland 17import { pipeline } from 'node:stream/promises';
f3253d0Takes: record, organize, and edit explainer videosJeremy Magland 18import ort from 'onnxruntime-node';
19import type { Take, Video } from '../shared/types.ts';
be26c78Update matte pipeline, server API and video pageJeremy Magland 20import { backgroundOf } from '../shared/types.ts';
f3253d0Takes: record, organize, and edit explainer videosJeremy Magland 21import { FPS } from '../shared/layout.ts';
22import { probe } from './media.ts';
307b4e2Run background removal on a remote GPU over sshJeremy Magland 23import { HttpError, exists, getConfig, getVideo, takeDir, updateVideo } from './store.ts';
f3253d0Takes: record, organize, and edit explainer videosJeremy Magland 24
25const MODEL_URL = 'https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_fp32.onnx';
26const MODEL_PATH = path.join(os.homedir(), '.cache', 'takes', 'models', 'rvm_mobilenetv3_fp32.onnx');
28interface Job {
29 videoId: string;
30 takeId: string;
be26c78Update matte pipeline, server API and video pageJeremy Magland 31 stage?: string; // what the running job is doing, for the UI
32 progress?: number; // fraction of the stage done, if known
be26c78Update matte pipeline, server API and video pageJeremy Magland 34 camera?: { width: number; height: number }; // size of the camera track, once running
35 height?: number; // lines being processed, once running
36 abort: AbortController;
37 restart?: boolean; // aborted because the processing resolution changed; run again
38 prev: Take['matte']; // persisted state before queueing, restored if a redo is cancelled
41const jobs = new Map<string, Job>(); // key: videoId/takeId, in queue order
42let pumping = false;
43let session: Promise<ort.InferenceSession> | null = null;
45const key = (videoId: string, takeId: string) => `${videoId}/${takeId}`;
47export const cutoutFile = (videoId: string, takeId: string) => path.join(takeDir(videoId, takeId), 'cutout.webm');
49// Merge the live queue state into a video before sending it to the client.
50export function withMatteStatus(v: Video): Video {
51 const takes = v.takes.map((t): Take => {
52 const j = jobs.get(key(v.id, t.id));
be26c78Update matte pipeline, server API and video pageJeremy Magland 53 return j ? { ...t, matte: { status: j.running ? 'running' : 'queued', stage: j.stage, progress: j.progress, updatedAt: new Date().toISOString() } } : t;
55 return { ...v, takes };
56}
be26c78Update matte pipeline, server API and video pageJeremy Magland 58const setMatte = (videoId: string, takeId: string, matte: Take['matte']) =>
59 updateVideo(videoId, (v) => {
60 const t = v.takes.find((t) => t.id === takeId);
61 if (t) t.matte = matte;
62 }).catch(() => {}); // the video or take may have been deleted meanwhile
64// The size a camera track is processed at: its own size, scaled down to at most `lines` lines.
65function processingSize(width: number, height: number, lines: number) {
66 const scale = Math.min(1, lines / height);
67 return { W: even(width * scale), H: even(height * scale), capped: scale < 1 };
68}
70// Queue every ready camera take of the video that has not been processed yet. With a take id, also
71// queue that take if it failed, was cancelled, or is done (to process it again at the current
72// resolution). A running job whose resolution no longer matches the setting is restarted.
73export async function requestMattes(videoId: string, opts: { takeId?: string } = {}) {
f3253d0Takes: record, organize, and edit explainer videosJeremy Magland 74 const v = await getVideo(videoId);
be26c78Update matte pipeline, server API and video pageJeremy Magland 75 const lines = backgroundOf(v.layout).resolution;
f3253d0Takes: record, organize, and edit explainer videosJeremy Magland 76 for (const t of v.takes) {
77 if (opts.takeId && t.id !== opts.takeId) continue;
be26c78Update matte pipeline, server API and video pageJeremy Magland 78 if (t.status !== 'ready' || !t.tracks.camera) continue;
79 const job = jobs.get(key(videoId, t.id));
80 if (job) {
81 if (job.running && job.camera && !job.restart && processingSize(job.camera.width, job.camera.height, lines).H !== job.height) {
82 job.restart = true;
83 job.abort.abort(new Error('Restarting at a new resolution'));
84 }
85 continue;
86 }
87 if (!opts.takeId) {
88 if (t.matte?.status === 'failed' || t.matte?.status === 'cancelled') continue;
89 if (t.matte?.status === 'done' && (await exists(cutoutFile(videoId, t.id)))) continue;
90 }
91 jobs.set(key(videoId, t.id), { videoId, takeId: t.id, running: false, abort: new AbortController(), prev: t.matte });
93 pump();
94}
be26c78Update matte pipeline, server API and video pageJeremy Magland 96// Cancel the queued and running jobs of a video (or of one take). A cancelled redo keeps its
97// previous cutout; otherwise the take is marked cancelled, so that it is not queued again until
98// asked for.
99export async function cancelMattes(videoId: string, takeId?: string) {
100 for (const job of [...jobs.values()]) {
101 if (job.videoId !== videoId || (takeId && job.takeId !== takeId)) continue;
102 job.restart = false;
103 if (job.running) job.abort.abort(new Error('Cancelled'));
104 else {
105 jobs.delete(key(job.videoId, job.takeId)); // before the await, so that the pump does not start it
106 await setMatte(job.videoId, job.takeId, cancelledState(job));
107 }
108 }
109}
111const cancelledState = (job: Job): Take['matte'] => (job.prev?.status === 'done' ? job.prev : { status: 'cancelled', updatedAt: new Date().toISOString() });
114 if (pumping) return;
115 pumping = true;
116 try {
117 for (;;) {
118 const job = jobs.values().next().value as Job | undefined;
119 if (!job) break;
120 job.running = true;
be26c78Update matte pipeline, server API and video pageJeremy Magland 121 const k = key(job.videoId, job.takeId);
123 const t0 = Date.now();
307b4e2Run background removal on a remote GPU over sshJeremy Magland 124 const { matteHost, matteSetup } = await getConfig();
be26c78Update matte pipeline, server API and video pageJeremy Magland 125 const { frames, H, capped } = await processTake(job, { host: matteHost, setup: matteSetup });
126 console.log(`matte: ${k}: ${frames} frames at ${H} lines in ${((Date.now() - t0) / 1000).toFixed(1)} s${matteHost ? ` on ${matteHost}` : ''}`);
127 await setMatte(job.videoId, job.takeId, { status: 'done', height: H, capped, updatedAt: new Date().toISOString() });
128 jobs.delete(k);
131 // Run it again next, with a fresh controller; it is still first in the queue.
132 console.log(`matte: ${k}: restarting at a new resolution`);
133 Object.assign(job, { running: false, stage: undefined, progress: undefined, camera: undefined, height: undefined, restart: false, abort: new AbortController() });
134 } else if (job.abort.signal.aborted) {
135 console.log(`matte: ${k}: cancelled`);
136 await setMatte(job.videoId, job.takeId, cancelledState(job));
137 jobs.delete(k);
138 } else {
139 console.warn(`matte failed for ${k}:`, (e as Error).message);
140 await setMatte(job.videoId, job.takeId, { status: 'failed', error: (e as Error).message, updatedAt: new Date().toISOString() });
141 jobs.delete(k);
142 }
144 }
145 } finally {
146 pumping = false;
147 }
148}
150async function ensureModel() {
151 if (await exists(MODEL_PATH)) return;
152 console.log(`Downloading background removal model to ${MODEL_PATH}`);
153 const res = await fetch(MODEL_URL);
154 if (!res.ok) throw new Error(`Could not download the model (${res.status}) from ${MODEL_URL}`);
155 await fs.mkdir(path.dirname(MODEL_PATH), { recursive: true });
156 const tmp = `${MODEL_PATH}.${process.pid}.tmp`;
157 await fs.writeFile(tmp, Buffer.from(await res.arrayBuffer()));
158 await fs.rename(tmp, MODEL_PATH);
159}
161function getSession() {
162 if (!session) {
163 session = ensureModel().then(() =>
164 ort.InferenceSession.create(MODEL_PATH, {
165 executionProviders: ['cpu'],
166 graphOptimizationLevel: 'all',
167 intraOpNumThreads: Math.max(1, os.cpus().length - 2),
168 }),
169 );
170 session.catch(() => (session = null)); // allow a retry after e.g. a failed download
171 }
172 return session;
173}
175// Yield fixed-size frames from a raw video stream. The same buffer is reused, so each frame must be
176// consumed before asking for the next.
177async function* readFrames(stream: Readable, size: number) {
178 const frame = Buffer.alloc(size);
179 let off = 0;
180 for await (const chunk of stream as AsyncIterable<Buffer>) {
181 let p = 0;
182 while (p < chunk.length) {
183 const n = Math.min(size - off, chunk.length - p);
184 chunk.copy(frame, off, p, p + n);
185 off += n;
186 p += n;
187 if (off === size) {
188 yield frame;
189 off = 0;
190 }
191 }
192 }
193}
195const even = (n: number) => Math.max(2, Math.round(n / 2) * 2);
307b4e2Run background removal on a remote GPU over sshJeremy Magland 197// ffmpeg arguments shared by the local and the remote path, so that both produce the same cutout.webm.
198const decoderArgs = (src: string, W: number, H: number) => ['-v', 'error', '-i', src, '-an', '-vf', `fps=${FPS}:start_time=0,scale=${W}:${H}`, '-f', 'rawvideo', '-pix_fmt', 'rgb24', '-'];
199const encoderArgs = (out: string, W: number, H: number) =>
200 ['-y', '-v', 'error', '-f', 'rawvideo', '-pix_fmt', 'rgba', '-s', `${W}x${H}`, '-r', String(FPS), '-i', '-',
201 '-c:v', 'libvpx-vp9', '-pix_fmt', 'yuva420p', '-deadline', 'realtime', '-cpu-used', '8', '-row-mt', '1', '-b:v', '3M', '-auto-alt-ref', '0', out];
203interface Plan {
204 src: string;
205 out: string;
206 W: number;
207 H: number;
208 ds: number;
209 total: number;
210}
be26c78Update matte pipeline, server API and video pageJeremy Magland 212async function processTake(job: Job, remote: { host: string; setup: string }) {
f3253d0Takes: record, organize, and edit explainer videosJeremy Magland 213 const src = path.join(takeDir(job.videoId, job.takeId), 'camera.webm');
214 const info = await probe(src);
215 if (!info.width || !info.height) throw new HttpError(400, 'Camera track has no video');
be26c78Update matte pipeline, server API and video pageJeremy Magland 216 const lines = backgroundOf((await getVideo(job.videoId)).layout).resolution;
217 const { W, H, capped } = processingSize(info.width, info.height, lines);
218 job.camera = { width: info.width, height: info.height };
219 job.height = H;
220 // Downsample ratio for the model's coarse pass, which takes most of the time. Following the RVM
221 // guidance it is at most 480 px wide (0.25 at 1080p, 0.375 at 720p); at low processing resolutions
222 // it is also capped at 0.6, so that a lower setting shrinks the coarse pass too (on a 12-thread
223 // laptop, 76 ms per frame at 426x240 instead of 250 ms at 1280x720).
224 const ds = Math.min(0.6, 480 / Math.max(W, H));
f3253d0Takes: record, organize, and edit explainer videosJeremy Magland 225 const total = Math.max(1, Math.round(info.duration * FPS));
227 const out = cutoutFile(job.videoId, job.takeId);
228 const tmp = `${out}.partial.webm`;
307b4e2Run background removal on a remote GPU over sshJeremy Magland 229 const plan: Plan = { src, out: tmp, W, H, ds, total };
230 try {
be26c78Update matte pipeline, server API and video pageJeremy Magland 231 job.abort.signal.throwIfAborted();
307b4e2Run background removal on a remote GPU over sshJeremy Magland 232 const n = remote.host ? await runRemote(job, plan, remote.host, remote.setup) : await runLocal(job, plan);
be26c78Update matte pipeline, server API and video pageJeremy Magland 233 job.abort.signal.throwIfAborted();
307b4e2Run background removal on a remote GPU over sshJeremy Magland 234 if (n === 0) throw new Error('No camera frames were decoded');
235 await fs.rename(tmp, out);
be26c78Update matte pipeline, server API and video pageJeremy Magland 236 return { frames: n, H, capped };
238 await fs.rm(tmp, { force: true });
239 throw e;
240 }
241}
969e06cSkip repeated and low-motion frames in background removal; add 180-line processingJeremy Magland 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.
254const SKIP_MOTION = 2;
256interface 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}
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).
265function 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}
290// Mean absolute difference of two RGB frames over a band (0 when the band is empty).
291function 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}
307b4e2Run background removal on a remote GPU over sshJeremy Magland 303async function runLocal(job: Job, { src, out, W, H, ds, total }: Plan): Promise<number> {
be26c78Update matte pipeline, server API and video pageJeremy Magland 304 job.stage = (await exists(MODEL_PATH)) ? 'Loading model' : 'Downloading model';
307b4e2Run background removal on a remote GPU over sshJeremy Magland 305 const sess = await getSession();
be26c78Update matte pipeline, server API and video pageJeremy Magland 306 job.stage = 'Removing background';
307 job.progress = 0;
308 const { signal } = job.abort;
309 signal.throwIfAborted();
307b4e2Run background removal on a remote GPU over sshJeremy Magland 310 const decoder = spawn('ffmpeg', decoderArgs(src, W, H), { stdio: ['ignore', 'pipe', 'pipe'] });
311 const encoder = spawn('ffmpeg', encoderArgs(out, W, H), { stdio: ['pipe', 'ignore', 'pipe'] });
be26c78Update matte pipeline, server API and video pageJeremy Magland 312 // SIGKILL, since ffmpeg does not act on SIGTERM while it is blocked reading its input.
313 const kill = () => {
314 encoder.stdin.destroy();
315 decoder.kill('SIGKILL');
316 encoder.kill('SIGKILL');
317 };
318 signal.addEventListener('abort', kill, { once: true });
319 encoder.stdin.on('error', () => {}); // e.g. EPIPE after a kill; the exit code reports failures
321 decoder.stderr.on('data', (d) => (stderr += d));
322 encoder.stderr.on('data', (d) => (stderr += d));
323 const decoderDone = once(decoder, 'close');
324 const encoderDone = once(encoder, 'close');
326 const plane = W * H;
327 const input = new Float32Array(3 * plane);
328 let rec: ort.Tensor[] = [0, 1, 2, 3].map(() => new ort.Tensor('float32', new Float32Array([0]), [1, 1, 1, 1]));
329 const dsr = new ort.Tensor('float32', new Float32Array([ds]), [1]);
969e06cSkip repeated and low-motion frames in background removal; add 180-line processingJeremy Magland 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 };
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 };
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;
387 for await (const frame of readFrames(decoder.stdout, plane * 3)) {
969e06cSkip repeated and low-motion frames in background removal; add 180-line processingJeremy Magland 389 if (prevRaw && frame.equals(prevRaw)) {
390 if (pending) pending.count++;
391 else await write(lastOut!);
392 continue;
969e06cSkip repeated and low-motion frames in background removal; add 180-line processingJeremy Magland 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 }
419 }
969e06cSkip repeated and low-motion frames in background removal; add 180-line processingJeremy Magland 421 if (pending) {
422 last = await matte(pending.raw);
423 processed++;
424 await write(rgbaOf(last), pending.count);
425 }
427 const [[dcode], [ecode]] = await Promise.all([decoderDone, encoderDone]);
428 if (dcode !== 0 || ecode !== 0) throw new Error(`ffmpeg failed: ${stderr.trim().split('\n').slice(-4).join('\n')}`);
969e06cSkip repeated and low-motion frames in background removal; add 180-line processingJeremy Magland 429 console.log(`matte: model ran on ${processed} of ${n} frames (${skipped} skipped for low motion, the rest repeated)`);
431 } catch (e) {
435 signal.removeEventListener('abort', kill);
437}
439// ---- remote processing over ssh ----
440//
be26c78Update matte pipeline, server API and video pageJeremy Magland 441// A single ssh connection carries the whole job: the camera track goes to the remote machine on stdin
442// (scaled down here first to the processing size, when that is smaller, so that far less has to be
443// uploaded),
307b4e2Run background removal on a remote GPU over sshJeremy Magland 444// the remote shell script below runs matte-worker.py on it in a temporary directory, and the
445// finished cutout.webm comes back on stdout. On first use the script creates a Python environment
446// with onnxruntime-gpu in ~/.cache/takes/venv on the remote machine; pip installs the CUDA and cuDNN
447// libraries as well, so the remote machine needs only an NVIDIA driver, Python 3.10 or later, and
448// ffmpeg (possibly made available by the setup command, e.g. "module load python ffmpeg").
be26c78Update matte pipeline, server API and video pageJeremy Magland 449// The remote side reports what it is doing on stderr with "takes:" lines (stage changes, frames
450// processed, the start of the upload, and the size of the result), which become the job's stage
451// and progress; other lines are logged. Cancelling a job closes the ssh connection; the remote worker then fails on its next progress
452// line (a broken pipe), and the script's trap removes the temporary directory.
307b4e2Run background removal on a remote GPU over sshJeremy Magland 453
454const WORKER = fs.readFile(new URL('./matte-worker.py', import.meta.url), 'utf8');
456const shq = (s: string) => `'${s.replaceAll("'", `'\\''`)}'`;
458function remoteScript(setup: string, worker: string) {
459 return `{ ${setup || ':'}
460} </dev/null
461set -e
be26c78Update matte pipeline, server API and video pageJeremy Magland 462echo "takes:stage Checking the Python environment on $(hostname)" >&2
464venv=$base/venv
465if [ ! -f "$venv/.ready" ]; then
466 python3 -c 'import sys; sys.exit(sys.version_info < (3, 10))' </dev/null || {
467 echo "Python 3.10 or later is needed on $(hostname), found: $(python3 --version 2>&1)" >&2; exit 1; }
be26c78Update matte pipeline, server API and video pageJeremy Magland 468 echo "takes:stage Installing onnxruntime-gpu on $(hostname) (first use only; about 2.5 GB)" >&2
470 mkdir -p "$base"
471 python3 -m venv "$venv" </dev/null
472 "$venv/bin/pip" install --quiet --disable-pip-version-check 'onnxruntime-gpu[cuda,cudnn]' numpy </dev/null >&2
473 touch "$venv/.ready"
474fi
475command -v ffmpeg >/dev/null || { echo "ffmpeg was not found on $(hostname)" >&2; exit 1; }
476work=$(mktemp -d)
477trap 'rm -rf "$work"' EXIT
478cat > "$work/worker.py" <<'TAKES_WORKER'
479${worker}
480TAKES_WORKER
483cd "$work"
484"$venv/bin/python" worker.py "$1" </dev/null >&2
be26c78Update matte pipeline, server API and video pageJeremy Magland 485echo "takes:result $(wc -c < out.webm)" >&2
487`;
488}
490async function runRemote(job: Job, { src, out, W, H, ds, total }: Plan, host: string, setup: string): Promise<number> {
491 const spec = { modelUrl: MODEL_URL, width: W, height: H, downsampleRatio: ds, decode: decoderArgs('in.webm', W, H), encode: encoderArgs('out.webm', W, H) };
492 const command = `bash -lc ${shq(remoteScript(setup, await WORKER))} takes-matte ${shq(JSON.stringify(spec))}`;
be26c78Update matte pipeline, server API and video pageJeremy Magland 493 const upload = await uploadSource(job, src, W, H, total / FPS, (f) => {
494 if (job.stage === 'Uploading camera track') job.progress = Math.min(1, f);
495 });
307b4e2Run background removal on a remote GPU over sshJeremy Magland 496 // BatchMode makes ssh fail instead of prompting for a password, since there is no terminal.
497 const ssh = spawn('ssh', ['-T', '-o', 'BatchMode=yes', '--', host, command], { stdio: ['pipe', 'pipe', 'pipe'] });
499 ssh.kill();
500 upload.kill();
501 };
502 job.abort.signal.addEventListener('abort', kill, { once: true });
503 job.stage = `Connecting to ${host}`;
504 job.progress = undefined;
505 let received = 0;
506 let resultSize = 0;
507 ssh.stdout.on('data', (c: Buffer) => {
508 received += c.length;
509 if (resultSize) job.progress = Math.min(1, received / resultSize);
510 });
307b4e2Run background removal on a remote GPU over sshJeremy Magland 511 const closed = new Promise<number | null>((resolve, reject) => {
512 ssh.on('error', reject);
513 ssh.on('close', resolve);
514 });
515 const lines: string[] = []; // recent diagnostics, for the error message
516 let n = 0;
517 let partial = '';
518 ssh.stderr.setEncoding('utf8');
519 ssh.stderr.on('data', (d: string) => {
520 const parts = (partial + d).split('\n');
521 partial = parts.pop()!;
522 for (const line of parts.map((l) => l.trimEnd())) {
be26c78Update matte pipeline, server API and video pageJeremy Magland 523 const m = /^takes:(progress|stage|upload|result) ?(.*)$/.exec(line);
524 if (m?.[1] === 'progress') {
525 n = parseInt(m[2]);
307b4e2Run background removal on a remote GPU over sshJeremy Magland 526 job.progress = Math.min(0.99, n / total);
be26c78Update matte pipeline, server API and video pageJeremy Magland 527 } else if (m?.[1] === 'stage') {
528 console.log(`matte (${host}): ${m[2]}`);
529 job.stage = m[2];
530 job.progress = m[2] === 'Removing background' ? 0 : undefined;
531 } else if (m?.[1] === 'upload') {
532 job.stage = 'Uploading camera track';
533 job.progress = Math.min(1, upload.fraction());
534 } else if (m?.[1] === 'result') {
535 resultSize = parseInt(m[2]) || 0;
536 job.stage = 'Receiving result';
537 job.progress = resultSize ? Math.min(1, received / resultSize) : undefined;
539 console.log(`matte (${host}): ${line}`);
540 lines.push(line);
541 if (lines.length > 20) lines.shift();
542 }
543 }
544 });
545 try {
546 const [code] = await Promise.all([
547 closed,
be26c78Update matte pipeline, server API and video pageJeremy Magland 548 pipeline(upload.stream, upload.counter, ssh.stdin).catch(() => {}), // an early exit shows up in the exit code
307b4e2Run background removal on a remote GPU over sshJeremy Magland 549 pipeline(ssh.stdout, createWriteStream(out)),
be26c78Update matte pipeline, server API and video pageJeremy Magland 550 // A failed downscale would otherwise go unnoticed and give a cutout of only part of the take.
551 upload.done.catch((e) => {
552 ssh.kill();
553 throw e;
554 }),
556 if (code !== 0) throw new Error(`Background removal on ${host} failed: ${lines.slice(-6).join('\n') || `ssh exited with code ${code}`}`);
be26c78Update matte pipeline, server API and video pageJeremy Magland 557 console.log(`matte (${host}): uploaded ${(upload.bytes() / 1e6).toFixed(1)} MB${upload.scaled ? ` (scaled to ${W}x${H})` : ''}`);
559 } catch (e) {
563 job.abort.signal.removeEventListener('abort', kill);
565}
be26c78Update matte pipeline, server API and video pageJeremy Magland 566
567// What is uploaded for a remote job: the camera track scaled down and re-encoded to the processing
568// size (H.264 at CRF 18, visually lossless: on real camera footage, 0.2 MB per 10 s at 240 lines
569// against 3.4 MB for the original 720p track), or the original
570// file when no scaling is needed. The progress is the fraction of the take sent so far; ffmpeg can
571// only run as far ahead of the upload as the pipe buffers allow.
572async function uploadSource(job: Job, src: string, W: number, H: number, duration: number, onProgress: (fraction: number) => void) {
573 let bytes = 0;
574 let fraction: () => number;
575 let stream: Readable;
576 let kill = () => {};
577 let done: Promise<void> = Promise.resolve();
578 const scaled = !!job.camera && H < job.camera.height;
579 if (scaled) {
580 const ff = spawn(
581 'ffmpeg',
582 ['-v', 'error', '-nostats', '-progress', 'pipe:2', '-i', src, '-an', '-vf', `fps=${FPS}:start_time=0,scale=${W}:${H}`,
583 '-c:v', 'libx264', '-preset', 'veryfast', '-crf', '18', '-pix_fmt', 'yuv420p', '-f', 'matroska', '-'],
584 { stdio: ['ignore', 'pipe', 'pipe'] },
585 );
586 let outTime = 0;
587 let errors = '';
588 ff.stderr.setEncoding('utf8');
589 ff.stderr.on('data', (d: string) => {
590 for (const line of d.split('\n')) {
591 const m = /^out_time_us=(\d+)/.exec(line);
592 if (m) outTime = parseInt(m[1]) / 1e6;
593 else if (line && !/^\w+=/.test(line)) errors += line + '\n';
594 }
595 onProgress(fraction());
596 });
597 fraction = () => (duration > 0 ? outTime / duration : 0);
598 stream = ff.stdout;
599 kill = () => ff.kill('SIGKILL');
600 done = once(ff, 'close').then(([code]) => {
601 if (code !== 0 && !job.abort.signal.aborted) throw new Error(`Could not scale the camera track for upload: ${errors.trim().split('\n').slice(-3).join('\n') || `ffmpeg exited with code ${code}`}`);
602 });
603 } else {
604 const size = (await fs.stat(src)).size;
605 fraction = () => bytes / size;
606 stream = createReadStream(src);
607 }
608 const counter = new Transform({
609 transform(chunk: Buffer, _enc, cb) {
610 bytes += chunk.length;
611 if (!scaled) onProgress(fraction());
612 cb(null, chunk);
613 },
614 });
615 return { stream, counter, scaled, done, kill, fraction: () => fraction(), bytes: () => bytes };
616}