Run background removal on a remote GPU over ssh
Settings gain an ssh host and a setup command. Each take's camera track is sent over one ssh connection, processed with onnxruntime-gpu in a temporary directory, and the cutout is copied back.
7 changed files+222−17
README.mdmodified+3−1View file
@@ -43,6 +43,8 @@ The library shows every video with its thumbnail, duration, and description. Dra
4343
4444 **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.
4545
46+**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 `camera.webm` over a single ssh connection, 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.
47+
4648 ## Storage layout
4749
4850 ```
@@ -68,6 +70,6 @@ Everything is plain files, so videos can be backed up, synced, or inspected with
6870 - Screen capture on Wayland depends on the browser's PipeWire support. Recent Chrome enables it by default; if the picker does not appear or the capture is black, check `chrome://flags/#enable-webrtc-pipewire-capturer`.
6971 - Only the microphone is recorded; system audio (for example, sound from a demo) is not captured.
7072 - The camera layout and background are set per video, not per segment.
71-- Background removal is a CPU-bound step and on a laptop it runs slower than real time; the matting quality also degrades with a cluttered background, poor lighting, or objects held close to the body, which the model may partly treat as background.
73+- Background removal on the CPU runs well below real time on a laptop (see above for running it on a remote GPU); the matting quality also degrades with a cluttered background, poor lighting, or objects held close to the body, which the model may partly treat as background.
7274 - Export speed depends on the CPU; with the `veryfast` x264 preset, a 1080p picture-in-picture export ran at roughly 3× real time on the development machine.
7375 - The server listens only on 127.0.0.1 and rejects requests addressed to other host names, but it has no authentication beyond that; it is meant for a single user on their own machine.
server/index.tsmodified+2−0View file
@@ -80,6 +80,8 @@ app.put('/api/config', json, async (req, res) => {
8080 if (typeof b.defaultCountdown === 'number') patch.defaultCountdown = Math.max(0, Math.min(10, Math.round(b.defaultCountdown)));
8181 if (b.defaultCameraHeight === 720 || b.defaultCameraHeight === 1080) patch.defaultCameraHeight = b.defaultCameraHeight;
8282 if (b.defaultOutputHeight === 720 || b.defaultOutputHeight === 1080) patch.defaultOutputHeight = b.defaultOutputHeight;
83+ if (typeof b.matteHost === 'string') patch.matteHost = b.matteHost.trim();
84+ if (typeof b.matteSetup === 'string') patch.matteSetup = b.matteSetup.trim();
8385 await updateConfig(patch);
8486 res.json({ ok: true });
8587 });
server/matte-worker.pyadded+75−0View file
@@ -0,0 +1,75 @@
1+# Background removal worker for a remote machine with an NVIDIA GPU. matte.ts copies this script to
2+# the remote machine together with the camera track and runs it in a temporary directory that holds
3+# in.webm. It runs the same Robust Video Matting model as the local path, through onnxruntime with
4+# the CUDA execution provider, and writes out.webm with the ffmpeg arguments given in the spec, so
5+# the result matches a local run. Progress is reported on stderr as "takes:progress <frames>" lines.
6+
7+import json
8+import os
9+import subprocess
10+import sys
11+import urllib.request
12+
13+import numpy as np
14+import onnxruntime as ort
15+
16+
17+def ensure_model(url):
18+ path = os.path.expanduser(os.path.join('~', '.cache', 'takes', 'models', os.path.basename(url)))
19+ if not os.path.exists(path):
20+ print(f'Downloading background removal model to {path}', file=sys.stderr, flush=True)
21+ os.makedirs(os.path.dirname(path), exist_ok=True)
22+ tmp = f'{path}.{os.getpid()}.tmp'
23+ urllib.request.urlretrieve(url, tmp)
24+ os.rename(tmp, path)
25+ return path
26+
27+
28+def to_u8(x):
29+ # Round half to even and clamp, as a Uint8ClampedArray does in the local path.
30+ return np.clip(np.rint(x * 255), 0, 255).astype(np.uint8)
31+
32+
33+def main():
34+ spec = json.loads(sys.argv[1])
35+ W, H = spec['width'], spec['height']
36+ model = ensure_model(spec['modelUrl'])
37+ if hasattr(ort, 'preload_dlls'):
38+ ort.preload_dlls() # the CUDA and cuDNN libraries installed by pip
39+ opts = ort.SessionOptions()
40+ opts.log_severity_level = 3 # errors only; a warning about memcpy nodes is expected for this model
41+ sess = ort.InferenceSession(model, opts, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
42+ if 'CUDAExecutionProvider' not in sess.get_providers():
43+ sys.exit('onnxruntime could not use the GPU (CUDAExecutionProvider is unavailable); see the messages above')
44+
45+ decoder = subprocess.Popen(['ffmpeg', *spec['decode']], stdout=subprocess.PIPE)
46+ encoder = subprocess.Popen(['ffmpeg', *spec['encode']], stdin=subprocess.PIPE)
47+ rec = [np.zeros((1, 1, 1, 1), np.float32)] * 4
48+ dsr = np.array([spec['downsampleRatio']], np.float32)
49+ rgba = np.empty((H, W, 4), np.uint8)
50+ size = W * H * 3
51+ n = 0
52+ while True:
53+ buf = decoder.stdout.read(size)
54+ if len(buf) < size:
55+ break
56+ frame = np.frombuffer(buf, np.uint8).reshape(H, W, 3)
57+ src = (frame.transpose(2, 0, 1)[None] / np.float32(255)).astype(np.float32)
58+ fgr, pha, *rec = sess.run(
59+ ['fgr', 'pha', 'r1o', 'r2o', 'r3o', 'r4o'],
60+ {'src': src, 'r1i': rec[0], 'r2i': rec[1], 'r3i': rec[2], 'r4i': rec[3], 'downsample_ratio': dsr},
61+ )
62+ rgba[..., :3] = to_u8(fgr[0].transpose(1, 2, 0))
63+ rgba[..., 3] = to_u8(pha[0, 0])
64+ encoder.stdin.write(rgba.tobytes())
65+ n += 1
66+ if n % 10 == 0:
67+ print(f'takes:progress {n}', file=sys.stderr, flush=True)
68+ encoder.stdin.close()
69+ if decoder.wait() != 0 or encoder.wait() != 0:
70+ sys.exit('ffmpeg failed')
71+ print(f'takes:progress {n}', file=sys.stderr, flush=True)
72+
73+
74+if __name__ == '__main__':
75+ main()
server/matte.tsmodified+123−16View file
@@ -1,20 +1,22 @@
11 // Background removal for camera tracks with Robust Video Matting (Lin et al., 2021), run on the CPU
2-// through onnxruntime. For each take, camera.webm is decoded to 30 fps RGB frames, the recurrent
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
33 // model predicts a foreground color and an alpha matte per frame, and the result is written as
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
77 import { spawn } from 'node:child_process';
88 import { once } from 'node:events';
9+import { createReadStream, createWriteStream } from 'node:fs';
910 import fs from 'node:fs/promises';
1011 import os from 'node:os';
1112 import path from 'node:path';
1213 import type { Readable } from 'node:stream';
14+import { pipeline } from 'node:stream/promises';
1315 import ort from 'onnxruntime-node';
1416 import type { Take, Video } from '../shared/types.ts';
1517 import { FPS } from '../shared/layout.ts';
1618 import { probe } from './media.ts';
17-import { HttpError, exists, getVideo, takeDir, updateVideo } from './store.ts';
19+import { HttpError, exists, getConfig, getVideo, takeDir, updateVideo } from './store.ts';
1820
1921 const MODEL_URL = 'https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_fp32.onnx';
2022 const MODEL_PATH = path.join(os.homedir(), '.cache', 'takes', 'models', 'rvm_mobilenetv3_fp32.onnx');
@@ -71,8 +73,9 @@ async function pump() {
7173 }).catch(() => {}); // the video or take may have been deleted meanwhile
7274 try {
7375 const t0 = Date.now();
74- const frames = await processTake(job);
75- console.log(`matte: ${job.videoId}/${job.takeId}: ${frames} frames in ${((Date.now() - t0) / 1000).toFixed(1)} s`);
76+ const { matteHost, matteSetup } = await getConfig();
77+ const frames = await processTake(job, { host: matteHost, setup: matteSetup });
78+ console.log(`matte: ${job.videoId}/${job.takeId}: ${frames} frames in ${((Date.now() - t0) / 1000).toFixed(1)} s${matteHost ? ` on ${matteHost}` : ''}`);
7679 await setMatte({ status: 'done', updatedAt: new Date().toISOString() });
7780 } catch (e) {
7881 console.warn(`matte failed for ${job.videoId}/${job.takeId}:`, (e as Error).message);
@@ -133,8 +136,22 @@ async function* readFrames(stream: Readable, size: number) {
133136
134137 const even = (n: number) => Math.max(2, Math.round(n / 2) * 2);
135138
136-async function processTake(job: Job): Promise<number> {
137- const sess = await getSession();
139+// ffmpeg arguments shared by the local and the remote path, so that both produce the same cutout.webm.
140+const 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', '-'];
141+const encoderArgs = (out: string, W: number, H: number) =>
142+ ['-y', '-v', 'error', '-f', 'rawvideo', '-pix_fmt', 'rgba', '-s', `${W}x${H}`, '-r', String(FPS), '-i', '-',
143+ '-c:v', 'libvpx-vp9', '-pix_fmt', 'yuva420p', '-deadline', 'realtime', '-cpu-used', '8', '-row-mt', '1', '-b:v', '3M', '-auto-alt-ref', '0', out];
144+
145+interface Plan {
146+ src: string;
147+ out: string;
148+ W: number;
149+ H: number;
150+ ds: number;
151+ total: number;
152+}
153+
154+async function processTake(job: Job, remote: { host: string; setup: string }): Promise<number> {
138155 const src = path.join(takeDir(job.videoId, job.takeId), 'camera.webm');
139156 const info = await probe(src);
140157 if (!info.width || !info.height) throw new HttpError(400, 'Camera track has no video');
@@ -148,13 +165,22 @@ async function processTake(job: Job): Promise<number> {
148165
149166 const out = cutoutFile(job.videoId, job.takeId);
150167 const tmp = `${out}.partial.webm`;
151- const decoder = spawn('ffmpeg', ['-v', 'error', '-i', src, '-an', '-vf', `fps=${FPS}:start_time=0,scale=${W}:${H}`, '-f', 'rawvideo', '-pix_fmt', 'rgb24', '-'], { stdio: ['ignore', 'pipe', 'pipe'] });
152- const encoder = spawn(
153- 'ffmpeg',
154- ['-y', '-v', 'error', '-f', 'rawvideo', '-pix_fmt', 'rgba', '-s', `${W}x${H}`, '-r', String(FPS), '-i', '-',
155- '-c:v', 'libvpx-vp9', '-pix_fmt', 'yuva420p', '-deadline', 'realtime', '-cpu-used', '8', '-row-mt', '1', '-b:v', '3M', '-auto-alt-ref', '0', tmp],
156- { stdio: ['pipe', 'ignore', 'pipe'] },
157- );
168+ const plan: Plan = { src, out: tmp, W, H, ds, total };
169+ try {
170+ const n = remote.host ? await runRemote(job, plan, remote.host, remote.setup) : await runLocal(job, plan);
171+ if (n === 0) throw new Error('No camera frames were decoded');
172+ await fs.rename(tmp, out);
173+ return n;
174+ } catch (e) {
175+ await fs.rm(tmp, { force: true });
176+ throw e;
177+ }
178+}
179+
180+async function runLocal(job: Job, { src, out, W, H, ds, total }: Plan): Promise<number> {
181+ const sess = await getSession();
182+ const decoder = spawn('ffmpeg', decoderArgs(src, W, H), { stdio: ['ignore', 'pipe', 'pipe'] });
183+ const encoder = spawn('ffmpeg', encoderArgs(out, W, H), { stdio: ['pipe', 'ignore', 'pipe'] });
158184 let stderr = '';
159185 decoder.stderr.on('data', (d) => (stderr += d));
160186 encoder.stderr.on('data', (d) => (stderr += d));
@@ -198,13 +224,94 @@ async function processTake(job: Job): Promise<number> {
198224 encoder.stdin.end();
199225 const [[dcode], [ecode]] = await Promise.all([decoderDone, encoderDone]);
200226 if (dcode !== 0 || ecode !== 0) throw new Error(`ffmpeg failed: ${stderr.trim().split('\n').slice(-4).join('\n')}`);
201- if (n === 0) throw new Error('No camera frames were decoded');
202- await fs.rename(tmp, out);
203227 return n;
204228 } catch (e) {
205229 decoder.kill();
206230 encoder.kill();
207- await fs.rm(tmp, { force: true });
231+ throw e;
232+ }
233+}
234+
235+// ---- remote processing over ssh ----
236+//
237+// A single ssh connection carries the whole job: camera.webm goes to the remote machine on stdin,
238+// the remote shell script below runs matte-worker.py on it in a temporary directory, and the
239+// finished cutout.webm comes back on stdout. On first use the script creates a Python environment
240+// with onnxruntime-gpu in ~/.cache/takes/venv on the remote machine; pip installs the CUDA and cuDNN
241+// libraries as well, so the remote machine needs only an NVIDIA driver, Python 3.10 or later, and
242+// ffmpeg (possibly made available by the setup command, e.g. "module load python ffmpeg").
243+
244+const WORKER = fs.readFile(new URL('./matte-worker.py', import.meta.url), 'utf8');
245+
246+const shq = (s: string) => `'${s.replaceAll("'", `'\\''`)}'`;
247+
248+function remoteScript(setup: string, worker: string) {
249+ return `{ ${setup || ':'}
250+} </dev/null
251+set -e
252+base=~/.cache/takes
253+venv=$base/venv
254+if [ ! -f "$venv/.ready" ]; then
255+ python3 -c 'import sys; sys.exit(sys.version_info < (3, 10))' </dev/null || {
256+ echo "Python 3.10 or later is needed on $(hostname), found: $(python3 --version 2>&1)" >&2; exit 1; }
257+ echo "Installing onnxruntime-gpu in $venv on $(hostname) (first use only; about 2.5 GB)" >&2
258+ rm -rf "$venv"
259+ mkdir -p "$base"
260+ python3 -m venv "$venv" </dev/null
261+ "$venv/bin/pip" install --quiet --disable-pip-version-check 'onnxruntime-gpu[cuda,cudnn]' numpy </dev/null >&2
262+ touch "$venv/.ready"
263+fi
264+command -v ffmpeg >/dev/null || { echo "ffmpeg was not found on $(hostname)" >&2; exit 1; }
265+work=$(mktemp -d)
266+trap 'rm -rf "$work"' EXIT
267+cat > "$work/worker.py" <<'TAKES_WORKER'
268+${worker}
269+TAKES_WORKER
270+cat > "$work/in.webm"
271+cd "$work"
272+"$venv/bin/python" worker.py "$1" </dev/null >&2
273+cat out.webm
274+`;
275+}
276+
277+async function runRemote(job: Job, { src, out, W, H, ds, total }: Plan, host: string, setup: string): Promise<number> {
278+ const spec = { modelUrl: MODEL_URL, width: W, height: H, downsampleRatio: ds, decode: decoderArgs('in.webm', W, H), encode: encoderArgs('out.webm', W, H) };
279+ const command = `bash -lc ${shq(remoteScript(setup, await WORKER))} takes-matte ${shq(JSON.stringify(spec))}`;
280+ // BatchMode makes ssh fail instead of prompting for a password, since there is no terminal.
281+ const ssh = spawn('ssh', ['-T', '-o', 'BatchMode=yes', '--', host, command], { stdio: ['pipe', 'pipe', 'pipe'] });
282+ const closed = new Promise<number | null>((resolve, reject) => {
283+ ssh.on('error', reject);
284+ ssh.on('close', resolve);
285+ });
286+ const lines: string[] = []; // recent diagnostics, for the error message
287+ let n = 0;
288+ let partial = '';
289+ ssh.stderr.setEncoding('utf8');
290+ ssh.stderr.on('data', (d: string) => {
291+ const parts = (partial + d).split('\n');
292+ partial = parts.pop()!;
293+ for (const line of parts.map((l) => l.trimEnd())) {
294+ const m = /^takes:progress (\d+)$/.exec(line);
295+ if (m) {
296+ n = parseInt(m[1]);
297+ job.progress = Math.min(0.99, n / total);
298+ } else if (line.trim()) {
299+ console.log(`matte (${host}): ${line}`);
300+ lines.push(line);
301+ if (lines.length > 20) lines.shift();
302+ }
303+ }
304+ });
305+ try {
306+ const [code] = await Promise.all([
307+ closed,
308+ pipeline(createReadStream(src), ssh.stdin).catch(() => {}), // an early exit shows up in the exit code
309+ pipeline(ssh.stdout, createWriteStream(out)),
310+ ]);
311+ if (code !== 0) throw new Error(`Background removal on ${host} failed: ${lines.slice(-6).join('\n') || `ssh exited with code ${code}`}`);
312+ return n;
313+ } catch (e) {
314+ ssh.kill();
208315 throw e;
209316 }
210317 }
server/store.tsmodified+2−0View file
@@ -85,6 +85,8 @@ const DEFAULT_CONFIG: AppConfig & { deepgramApiKey: string } = {
8585 defaultCountdown: 3,
8686 defaultCameraHeight: 720,
8787 defaultOutputHeight: 1080,
88+ matteHost: '',
89+ matteSetup: '',
8890 };
8991
9092 export type FullConfig = typeof DEFAULT_CONFIG;
shared/types.tsmodified+2−0View file
@@ -106,6 +106,8 @@ export interface AppConfig {
106106 defaultCountdown: number;
107107 defaultCameraHeight: 720 | 1080;
108108 defaultOutputHeight: 720 | 1080;
109+ matteHost: string; // ssh destination for background removal on a GPU; empty for this computer
110+ matteSetup: string; // shell commands run on that host first, e.g. "module load python ffmpeg"
109111 }
110112
111113 export interface PublicConfig extends AppConfig {
src/pages/Settings.tsxmodified+15−0View file
@@ -80,6 +80,21 @@ export function Settings() {
8080 </label>
8181 </section>
8282
83+ <section className="panel">
84+ <h4>Background removal</h4>
85+ <p className="muted small">
86+ By default the camera background is removed on this computer’s CPU. To use a machine with an NVIDIA GPU instead, enter its ssh destination (for example <code>office</code> or <code>user@host</code>). It must accept key-based login without a password prompt and have Python 3.10 or later and ffmpeg; the first job installs onnxruntime-gpu in <code>~/.cache/takes/venv</code> there. Each camera track is copied to that machine for processing and then deleted from it.
87+ </p>
88+ <label className="field">
89+ <span>ssh host</span>
90+ <input placeholder="This computer" defaultValue={cfg.matteHost} onBlur={(e) => e.target.value.trim() !== cfg.matteHost && save({ matteHost: e.target.value })} />
91+ </label>
92+ <label className="field">
93+ <span>Setup command on that host</span>
94+ <input placeholder="e.g. module load python ffmpeg" defaultValue={cfg.matteSetup} onBlur={(e) => e.target.value.trim() !== cfg.matteSetup && save({ matteSetup: e.target.value })} />
95+ </label>
96+ </section>
97+
8398 <section className="panel">
8499 <h4>Defaults</h4>
85100 <label className="row">