server/matte.ts at 6ec86a84f34da00ad6d472efdc039cc5caf81f93

1// Background removal for camera tracks with Robust Video Matting (Lin et al., 2021), run on the CPU

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

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.

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

7// or running job can be cancelled.

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

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

11import { createReadStream, createWriteStream } from 'node:fs';

12import fs from 'node:fs/promises';

13import os from 'node:os';

14import path from 'node:path';

15import { Transform, type Readable } from 'node:stream';

16import { pipeline } from 'node:stream/promises';

17import ort from 'onnxruntime-node';

18import type { Take, Video } from '../shared/types.ts';

19import { backgroundOf } from '../shared/types.ts';

20import { FPS } from '../shared/layout.ts';

21import { probe } from './media.ts';

22import { HttpError, exists, getConfig, getVideo, takeDir, updateVideo } from './store.ts';

24const MODEL_URL = 'https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_fp32.onnx';

25const MODEL_PATH = path.join(os.homedir(), '.cache', 'takes', 'models', 'rvm_mobilenetv3_fp32.onnx');

27interface Job {

28 videoId: string;

29 takeId: string;

30 stage?: string; // what the running job is doing, for the UI

31 progress?: number; // fraction of the stage done, if known

32 running: boolean;

33 camera?: { width: number; height: number }; // size of the camera track, once running

34 height?: number; // lines being processed, once running

35 abort: AbortController;

36 restart?: boolean; // aborted because the processing resolution changed; run again

37 prev: Take['matte']; // persisted state before queueing, restored if a redo is cancelled

38}

40const jobs = new Map<string, Job>(); // key: videoId/takeId, in queue order

41let pumping = false;

42let session: Promise<ort.InferenceSession> | null = null;

44const key = (videoId: string, takeId: string) => `${videoId}/${takeId}`;

46export const cutoutFile = (videoId: string, takeId: string) => path.join(takeDir(videoId, takeId), 'cutout.webm');

48// Merge the live queue state into a video before sending it to the client.

49export function withMatteStatus(v: Video): Video {

50 const takes = v.takes.map((t): Take => {

51 const j = jobs.get(key(v.id, t.id));

52 return j ? { ...t, matte: { status: j.running ? 'running' : 'queued', stage: j.stage, progress: j.progress, updatedAt: new Date().toISOString() } } : t;

53 });

54 return { ...v, takes };

55}

57const setMatte = (videoId: string, takeId: string, matte: Take['matte']) =>

58 updateVideo(videoId, (v) => {

59 const t = v.takes.find((t) => t.id === takeId);

60 if (t) t.matte = matte;

61 }).catch(() => {}); // the video or take may have been deleted meanwhile

63// The size a camera track is processed at: its own size, scaled down to at most `lines` lines.

64function processingSize(width: number, height: number, lines: number) {

65 const scale = Math.min(1, lines / height);

66 return { W: even(width * scale), H: even(height * scale), capped: scale < 1 };

67}

69// Queue every ready camera take of the video that has not been processed yet. With a take id, also

70// queue that take if it failed, was cancelled, or is done (to process it again at the current

71// resolution). A running job whose resolution no longer matches the setting is restarted.

72export async function requestMattes(videoId: string, opts: { takeId?: string } = {}) {

73 const v = await getVideo(videoId);

74 const lines = backgroundOf(v.layout).resolution;

75 for (const t of v.takes) {

76 if (opts.takeId && t.id !== opts.takeId) continue;

77 if (t.status !== 'ready' || !t.tracks.camera) continue;

78 const job = jobs.get(key(videoId, t.id));

79 if (job) {

80 if (job.running && job.camera && !job.restart && processingSize(job.camera.width, job.camera.height, lines).H !== job.height) {

81 job.restart = true;

82 job.abort.abort(new Error('Restarting at a new resolution'));

83 }

84 continue;

85 }

86 if (!opts.takeId) {

87 if (t.matte?.status === 'failed' || t.matte?.status === 'cancelled') continue;

88 if (t.matte?.status === 'done' && (await exists(cutoutFile(videoId, t.id)))) continue;

89 }

90 jobs.set(key(videoId, t.id), { videoId, takeId: t.id, running: false, abort: new AbortController(), prev: t.matte });

91 }

92 pump();

93}

95// Cancel the queued and running jobs of a video (or of one take). A cancelled redo keeps its

96// previous cutout; otherwise the take is marked cancelled, so that it is not queued again until

97// asked for.

98export async function cancelMattes(videoId: string, takeId?: string) {

99 for (const job of [...jobs.values()]) {

100 if (job.videoId !== videoId || (takeId && job.takeId !== takeId)) continue;

101 job.restart = false;

102 if (job.running) job.abort.abort(new Error('Cancelled'));

103 else {

104 jobs.delete(key(job.videoId, job.takeId)); // before the await, so that the pump does not start it

105 await setMatte(job.videoId, job.takeId, cancelledState(job));

106 }

107 }

108}

109

110const cancelledState = (job: Job): Take['matte'] => (job.prev?.status === 'done' ? job.prev : { status: 'cancelled', updatedAt: new Date().toISOString() });

111

112async function pump() {

113 if (pumping) return;

114 pumping = true;

115 try {

116 for (;;) {

117 const job = jobs.values().next().value as Job | undefined;

118 if (!job) break;

119 job.running = true;

120 const k = key(job.videoId, job.takeId);

121 try {

122 const t0 = Date.now();

123 const { matteHost, matteSetup } = await getConfig();

124 const { frames, H, capped } = await processTake(job, { host: matteHost, setup: matteSetup });

125 console.log(`matte: ${k}: ${frames} frames at ${H} lines in ${((Date.now() - t0) / 1000).toFixed(1)} s${matteHost ? ` on ${matteHost}` : ''}`);

126 await setMatte(job.videoId, job.takeId, { status: 'done', height: H, capped, updatedAt: new Date().toISOString() });

127 jobs.delete(k);

128 } catch (e) {

129 if (job.restart) {

130 // Run it again next, with a fresh controller; it is still first in the queue.

131 console.log(`matte: ${k}: restarting at a new resolution`);

132 Object.assign(job, { running: false, stage: undefined, progress: undefined, camera: undefined, height: undefined, restart: false, abort: new AbortController() });

133 } else if (job.abort.signal.aborted) {

134 console.log(`matte: ${k}: cancelled`);

135 await setMatte(job.videoId, job.takeId, cancelledState(job));

136 jobs.delete(k);

137 } else {

138 console.warn(`matte failed for ${k}:`, (e as Error).message);

139 await setMatte(job.videoId, job.takeId, { status: 'failed', error: (e as Error).message, updatedAt: new Date().toISOString() });

140 jobs.delete(k);

141 }

142 }

143 }

144 } finally {

145 pumping = false;

146 }

147}

148

149async function ensureModel() {

150 if (await exists(MODEL_PATH)) return;

151 console.log(`Downloading background removal model to ${MODEL_PATH}`);

152 const res = await fetch(MODEL_URL);

153 if (!res.ok) throw new Error(`Could not download the model (${res.status}) from ${MODEL_URL}`);

154 await fs.mkdir(path.dirname(MODEL_PATH), { recursive: true });

155 const tmp = `${MODEL_PATH}.${process.pid}.tmp`;

156 await fs.writeFile(tmp, Buffer.from(await res.arrayBuffer()));

157 await fs.rename(tmp, MODEL_PATH);

158}

159

160function getSession() {

161 if (!session) {

162 session = ensureModel().then(() =>

163 ort.InferenceSession.create(MODEL_PATH, {

164 executionProviders: ['cpu'],

165 graphOptimizationLevel: 'all',

166 intraOpNumThreads: Math.max(1, os.cpus().length - 2),

167 }),

168 );

169 session.catch(() => (session = null)); // allow a retry after e.g. a failed download

170 }

171 return session;

172}

173

174// Yield fixed-size frames from a raw video stream. The same buffer is reused, so each frame must be

175// consumed before asking for the next.

176async function* readFrames(stream: Readable, size: number) {

177 const frame = Buffer.alloc(size);

178 let off = 0;

179 for await (const chunk of stream as AsyncIterable<Buffer>) {

180 let p = 0;

181 while (p < chunk.length) {

182 const n = Math.min(size - off, chunk.length - p);

183 chunk.copy(frame, off, p, p + n);

184 off += n;

185 p += n;

186 if (off === size) {

187 yield frame;

188 off = 0;

189 }

190 }

191 }

192}

193

194const even = (n: number) => Math.max(2, Math.round(n / 2) * 2);

195

196// ffmpeg arguments shared by the local and the remote path, so that both produce the same cutout.webm.

197const 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', '-'];

198const encoderArgs = (out: string, W: number, H: number) =>

199 ['-y', '-v', 'error', '-f', 'rawvideo', '-pix_fmt', 'rgba', '-s', `${W}x${H}`, '-r', String(FPS), '-i', '-',

200 '-c:v', 'libvpx-vp9', '-pix_fmt', 'yuva420p', '-deadline', 'realtime', '-cpu-used', '8', '-row-mt', '1', '-b:v', '3M', '-auto-alt-ref', '0', out];

201

202interface Plan {

203 src: string;

204 out: string;

205 W: number;

206 H: number;

207 ds: number;

208 total: number;

209}

210

211async function processTake(job: Job, remote: { host: string; setup: string }) {

212 const src = path.join(takeDir(job.videoId, job.takeId), 'camera.webm');

213 const info = await probe(src);

214 if (!info.width || !info.height) throw new HttpError(400, 'Camera track has no video');

215 const lines = backgroundOf((await getVideo(job.videoId)).layout).resolution;

216 const { W, H, capped } = processingSize(info.width, info.height, lines);

217 job.camera = { width: info.width, height: info.height };

218 job.height = H;

219 // Downsample ratio for the model's coarse pass, which takes most of the time. Following the RVM

220 // guidance it is at most 480 px wide (0.25 at 1080p, 0.375 at 720p); at low processing resolutions

221 // it is also capped at 0.6, so that a lower setting shrinks the coarse pass too (on a 12-thread

222 // laptop, 76 ms per frame at 426x240 instead of 250 ms at 1280x720).

223 const ds = Math.min(0.6, 480 / Math.max(W, H));

224 const total = Math.max(1, Math.round(info.duration * FPS));

225

226 const out = cutoutFile(job.videoId, job.takeId);

227 const tmp = `${out}.partial.webm`;

228 const plan: Plan = { src, out: tmp, W, H, ds, total };

229 try {

230 job.abort.signal.throwIfAborted();

231 const n = remote.host ? await runRemote(job, plan, remote.host, remote.setup) : await runLocal(job, plan);

232 job.abort.signal.throwIfAborted();

233 if (n === 0) throw new Error('No camera frames were decoded');

234 await fs.rename(tmp, out);

235 return { frames: n, H, capped };

236 } catch (e) {

237 await fs.rm(tmp, { force: true });

238 throw e;

239 }

240}

241

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

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

244 const sess = await getSession();

245 job.stage = 'Removing background';

246 job.progress = 0;

247 const { signal } = job.abort;

248 signal.throwIfAborted();

249 const decoder = spawn('ffmpeg', decoderArgs(src, W, H), { stdio: ['ignore', 'pipe', 'pipe'] });

250 const encoder = spawn('ffmpeg', encoderArgs(out, W, H), { stdio: ['pipe', 'ignore', 'pipe'] });

251 // SIGKILL, since ffmpeg does not act on SIGTERM while it is blocked reading its input.

252 const kill = () => {

253 encoder.stdin.destroy();

254 decoder.kill('SIGKILL');

255 encoder.kill('SIGKILL');

256 };

257 signal.addEventListener('abort', kill, { once: true });

258 encoder.stdin.on('error', () => {}); // e.g. EPIPE after a kill; the exit code reports failures

259 let stderr = '';

260 decoder.stderr.on('data', (d) => (stderr += d));

261 encoder.stderr.on('data', (d) => (stderr += d));

262 const decoderDone = once(decoder, 'close');

263 const encoderDone = once(encoder, 'close');

264

265 const plane = W * H;

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

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

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

269 let n = 0;

270 try {

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

272 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;

277 }

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;

295 }

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);

299 }

300 signal.throwIfAborted();

301 encoder.stdin.end();

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

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

304 return n;

305 } catch (e) {

306 kill();

307 throw e;

308 } finally {

309 signal.removeEventListener('abort', kill);

310 }

311}

312

313// ---- remote processing over ssh ----

314//

315// A single ssh connection carries the whole job: the camera track goes to the remote machine on stdin

316// (scaled down here first to the processing size, when that is smaller, so that far less has to be

317// uploaded),

318// the remote shell script below runs matte-worker.py on it in a temporary directory, and the

319// finished cutout.webm comes back on stdout. On first use the script creates a Python environment

320// with onnxruntime-gpu in ~/.cache/takes/venv on the remote machine; pip installs the CUDA and cuDNN

321// libraries as well, so the remote machine needs only an NVIDIA driver, Python 3.10 or later, and

322// ffmpeg (possibly made available by the setup command, e.g. "module load python ffmpeg").

323// The remote side reports what it is doing on stderr with "takes:" lines (stage changes, frames

324// processed, the start of the upload, and the size of the result), which become the job's stage

325// and progress; other lines are logged. Cancelling a job closes the ssh connection; the remote worker then fails on its next progress

326// line (a broken pipe), and the script's trap removes the temporary directory.

327

328const WORKER = fs.readFile(new URL('./matte-worker.py', import.meta.url), 'utf8');

329

330const shq = (s: string) => `'${s.replaceAll("'", `'\\''`)}'`;

331

332function remoteScript(setup: string, worker: string) {

333 return `{ ${setup || ':'}

334} </dev/null

335set -e

336echo "takes:stage Checking the Python environment on $(hostname)" >&2

337base=~/.cache/takes

338venv=$base/venv

339if [ ! -f "$venv/.ready" ]; then

340 python3 -c 'import sys; sys.exit(sys.version_info < (3, 10))' </dev/null || {

341 echo "Python 3.10 or later is needed on $(hostname), found: $(python3 --version 2>&1)" >&2; exit 1; }

342 echo "takes:stage Installing onnxruntime-gpu on $(hostname) (first use only; about 2.5 GB)" >&2

343 rm -rf "$venv"

344 mkdir -p "$base"

345 python3 -m venv "$venv" </dev/null

346 "$venv/bin/pip" install --quiet --disable-pip-version-check 'onnxruntime-gpu[cuda,cudnn]' numpy </dev/null >&2

347 touch "$venv/.ready"

348fi

349command -v ffmpeg >/dev/null || { echo "ffmpeg was not found on $(hostname)" >&2; exit 1; }

350work=$(mktemp -d)

351trap 'rm -rf "$work"' EXIT

352cat > "$work/worker.py" <<'TAKES_WORKER'

353${worker}

354TAKES_WORKER

355echo "takes:upload" >&2

356cat > "$work/in.webm"

357cd "$work"

358"$venv/bin/python" worker.py "$1" </dev/null >&2

359echo "takes:result $(wc -c < out.webm)" >&2

360cat out.webm

361`;

362}

363

364async function runRemote(job: Job, { src, out, W, H, ds, total }: Plan, host: string, setup: string): Promise<number> {

365 const spec = { modelUrl: MODEL_URL, width: W, height: H, downsampleRatio: ds, decode: decoderArgs('in.webm', W, H), encode: encoderArgs('out.webm', W, H) };

366 const command = `bash -lc ${shq(remoteScript(setup, await WORKER))} takes-matte ${shq(JSON.stringify(spec))}`;

367 const upload = await uploadSource(job, src, W, H, total / FPS, (f) => {

368 if (job.stage === 'Uploading camera track') job.progress = Math.min(1, f);

369 });

370 // BatchMode makes ssh fail instead of prompting for a password, since there is no terminal.

371 const ssh = spawn('ssh', ['-T', '-o', 'BatchMode=yes', '--', host, command], { stdio: ['pipe', 'pipe', 'pipe'] });

372 const kill = () => {

373 ssh.kill();

374 upload.kill();

375 };

376 job.abort.signal.addEventListener('abort', kill, { once: true });

377 job.stage = `Connecting to ${host}`;

378 job.progress = undefined;

379 let received = 0;

380 let resultSize = 0;

381 ssh.stdout.on('data', (c: Buffer) => {

382 received += c.length;

383 if (resultSize) job.progress = Math.min(1, received / resultSize);

384 });

385 const closed = new Promise<number | null>((resolve, reject) => {

386 ssh.on('error', reject);

387 ssh.on('close', resolve);

388 });

389 const lines: string[] = []; // recent diagnostics, for the error message

390 let n = 0;

391 let partial = '';

392 ssh.stderr.setEncoding('utf8');

393 ssh.stderr.on('data', (d: string) => {

394 const parts = (partial + d).split('\n');

395 partial = parts.pop()!;

396 for (const line of parts.map((l) => l.trimEnd())) {

397 const m = /^takes:(progress|stage|upload|result) ?(.*)$/.exec(line);

398 if (m?.[1] === 'progress') {

399 n = parseInt(m[2]);

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

401 } else if (m?.[1] === 'stage') {

402 console.log(`matte (${host}): ${m[2]}`);

403 job.stage = m[2];

404 job.progress = m[2] === 'Removing background' ? 0 : undefined;

405 } else if (m?.[1] === 'upload') {

406 job.stage = 'Uploading camera track';

407 job.progress = Math.min(1, upload.fraction());

408 } else if (m?.[1] === 'result') {

409 resultSize = parseInt(m[2]) || 0;

410 job.stage = 'Receiving result';

411 job.progress = resultSize ? Math.min(1, received / resultSize) : undefined;

412 } else if (line.trim()) {

413 console.log(`matte (${host}): ${line}`);

414 lines.push(line);

415 if (lines.length > 20) lines.shift();

416 }

417 }

418 });

419 try {

420 const [code] = await Promise.all([

421 closed,

422 pipeline(upload.stream, upload.counter, ssh.stdin).catch(() => {}), // an early exit shows up in the exit code

423 pipeline(ssh.stdout, createWriteStream(out)),

424 // A failed downscale would otherwise go unnoticed and give a cutout of only part of the take.

425 upload.done.catch((e) => {

426 ssh.kill();

427 throw e;

428 }),

429 ]);

430 if (code !== 0) throw new Error(`Background removal on ${host} failed: ${lines.slice(-6).join('\n') || `ssh exited with code ${code}`}`);

431 console.log(`matte (${host}): uploaded ${(upload.bytes() / 1e6).toFixed(1)} MB${upload.scaled ? ` (scaled to ${W}x${H})` : ''}`);

432 return n;

433 } catch (e) {

434 kill();

435 throw e;

436 } finally {

437 job.abort.signal.removeEventListener('abort', kill);

438 }

439}

440

441// What is uploaded for a remote job: the camera track scaled down and re-encoded to the processing

442// size (H.264 at CRF 18, visually lossless: on real camera footage, 0.2 MB per 10 s at 240 lines

443// against 3.4 MB for the original 720p track), or the original

444// file when no scaling is needed. The progress is the fraction of the take sent so far; ffmpeg can

445// only run as far ahead of the upload as the pipe buffers allow.

446async function uploadSource(job: Job, src: string, W: number, H: number, duration: number, onProgress: (fraction: number) => void) {

447 let bytes = 0;

448 let fraction: () => number;

449 let stream: Readable;

450 let kill = () => {};

451 let done: Promise<void> = Promise.resolve();

452 const scaled = !!job.camera && H < job.camera.height;

453 if (scaled) {

454 const ff = spawn(

455 'ffmpeg',

456 ['-v', 'error', '-nostats', '-progress', 'pipe:2', '-i', src, '-an', '-vf', `fps=${FPS}:start_time=0,scale=${W}:${H}`,

457 '-c:v', 'libx264', '-preset', 'veryfast', '-crf', '18', '-pix_fmt', 'yuv420p', '-f', 'matroska', '-'],

458 { stdio: ['ignore', 'pipe', 'pipe'] },

459 );

460 let outTime = 0;

461 let errors = '';

462 ff.stderr.setEncoding('utf8');

463 ff.stderr.on('data', (d: string) => {

464 for (const line of d.split('\n')) {

465 const m = /^out_time_us=(\d+)/.exec(line);

466 if (m) outTime = parseInt(m[1]) / 1e6;

467 else if (line && !/^\w+=/.test(line)) errors += line + '\n';

468 }

469 onProgress(fraction());

470 });

471 fraction = () => (duration > 0 ? outTime / duration : 0);

472 stream = ff.stdout;

473 kill = () => ff.kill('SIGKILL');

474 done = once(ff, 'close').then(([code]) => {

475 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}`}`);

476 });

477 } else {

478 const size = (await fs.stat(src)).size;

479 fraction = () => bytes / size;

480 stream = createReadStream(src);

481 }

482 const counter = new Transform({

483 transform(chunk: Buffer, _enc, cb) {

484 bytes += chunk.length;

485 if (!scaled) onProgress(fraction());

486 cb(null, chunk);

487 },

488 });

489 return { stream, counter, scaled, done, kill, fraction: () => fraction(), bytes: () => bytes };

490}