1// 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
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.
7import { spawn } from 'node:child_process';
8import { once } from 'node:events';
9import fs from 'node:fs/promises';
10import os from 'node:os';
11import path from 'node:path';
12import type { Readable } from 'node:stream';
13import ort from 'onnxruntime-node';
14import type { Take, Video } from '../shared/types.ts';
15import { FPS } from '../shared/layout.ts';
16import { probe } from './media.ts';
17import { HttpError, exists, getVideo, takeDir, updateVideo } from './store.ts';
19const MODEL_URL = 'https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_fp32.onnx';
20const MODEL_PATH = path.join(os.homedir(), '.cache', 'takes', 'models', 'rvm_mobilenetv3_fp32.onnx');
22interface Job {
23 videoId: string;
24 takeId: string;
25 progress: number;
26 running: boolean;
27}
29const jobs = new Map<string, Job>(); // key: videoId/takeId, in queue order
30let pumping = false;
31let session: Promise<ort.InferenceSession> | null = null;
33const key = (videoId: string, takeId: string) => `${videoId}/${takeId}`;
35export const cutoutFile = (videoId: string, takeId: string) => path.join(takeDir(videoId, takeId), 'cutout.webm');
37// Merge the live queue state into a video before sending it to the client.
38export function withMatteStatus(v: Video): Video {
39 const takes = v.takes.map((t): Take => {
40 const j = jobs.get(key(v.id, t.id));
41 return j ? { ...t, matte: { status: j.running ? 'running' : 'queued', progress: j.progress, updatedAt: new Date().toISOString() } } : t;
42 });
43 return { ...v, takes };
44}
46// Queue every ready camera take of the video that has not been processed yet (or failed, if retry).
47export async function requestMattes(videoId: string, opts: { retry?: boolean; takeId?: string } = {}) {
48 const v = await getVideo(videoId);
49 for (const t of v.takes) {
50 if (opts.takeId && t.id !== opts.takeId) continue;
51 if (t.status !== 'ready' || !t.tracks.camera || jobs.has(key(videoId, t.id))) continue;
52 if (t.matte?.status === 'done' && (await exists(cutoutFile(videoId, t.id)))) continue;
53 if (t.matte?.status === 'failed' && !opts.retry) continue;
54 jobs.set(key(videoId, t.id), { videoId, takeId: t.id, progress: 0, running: false });
55 }
56 pump();
57}
59async function pump() {
60 if (pumping) return;
61 pumping = true;
62 try {
63 for (;;) {
64 const job = jobs.values().next().value as Job | undefined;
65 if (!job) break;
66 job.running = true;
67 const setMatte = (matte: Take['matte']) =>
68 updateVideo(job.videoId, (v) => {
69 const t = v.takes.find((t) => t.id === job.takeId);
70 if (t) t.matte = matte;
71 }).catch(() => {}); // the video or take may have been deleted meanwhile
72 try {
73 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 await setMatte({ status: 'done', updatedAt: new Date().toISOString() });
77 } catch (e) {
78 console.warn(`matte failed for ${job.videoId}/${job.takeId}:`, (e as Error).message);
79 await setMatte({ status: 'failed', error: (e as Error).message, updatedAt: new Date().toISOString() });
80 } finally {
81 jobs.delete(key(job.videoId, job.takeId));
82 }
83 }
84 } finally {
85 pumping = false;
86 }
87}
89async function ensureModel() {
90 if (await exists(MODEL_PATH)) return;
91 console.log(`Downloading background removal model to ${MODEL_PATH}`);
92 const res = await fetch(MODEL_URL);
93 if (!res.ok) throw new Error(`Could not download the model (${res.status}) from ${MODEL_URL}`);
94 await fs.mkdir(path.dirname(MODEL_PATH), { recursive: true });
95 const tmp = `${MODEL_PATH}.${process.pid}.tmp`;
96 await fs.writeFile(tmp, Buffer.from(await res.arrayBuffer()));
97 await fs.rename(tmp, MODEL_PATH);
98}
100function getSession() {
101 if (!session) {
102 session = ensureModel().then(() =>
103 ort.InferenceSession.create(MODEL_PATH, {
104 executionProviders: ['cpu'],
105 graphOptimizationLevel: 'all',
106 intraOpNumThreads: Math.max(1, os.cpus().length - 2),
107 }),
108 );
109 session.catch(() => (session = null)); // allow a retry after e.g. a failed download
110 }
111 return session;
112}
114// Yield fixed-size frames from a raw video stream. The same buffer is reused, so each frame must be
115// consumed before asking for the next.
116async function* readFrames(stream: Readable, size: number) {
117 const frame = Buffer.alloc(size);
118 let off = 0;
119 for await (const chunk of stream as AsyncIterable<Buffer>) {
120 let p = 0;
121 while (p < chunk.length) {
122 const n = Math.min(size - off, chunk.length - p);
123 chunk.copy(frame, off, p, p + n);
124 off += n;
125 p += n;
126 if (off === size) {
127 yield frame;
128 off = 0;
129 }
130 }
131 }
132}
134const even = (n: number) => Math.max(2, Math.round(n / 2) * 2);
136async function processTake(job: Job): Promise<number> {
137 const sess = await getSession();
138 const src = path.join(takeDir(job.videoId, job.takeId), 'camera.webm');
139 const info = await probe(src);
140 if (!info.width || !info.height) throw new HttpError(400, 'Camera track has no video');
141 // Process at the camera's resolution, capped at 1080 lines.
142 const scale = Math.min(1, 1080 / info.height);
143 const W = even(info.width * scale);
144 const H = even(info.height * scale);
145 // Downsample ratio for the model's coarse pass, following the RVM guidance (0.25 at 1080p, 0.375 at 720p).
146 const ds = Math.min(1, 480 / Math.max(W, H));
147 const total = Math.max(1, Math.round(info.duration * FPS));
149 const out = cutoutFile(job.videoId, job.takeId);
150 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 );
158 let stderr = '';
159 decoder.stderr.on('data', (d) => (stderr += d));
160 encoder.stderr.on('data', (d) => (stderr += d));
161 const decoderDone = once(decoder, 'close');
162 const encoderDone = once(encoder, 'close');
164 const plane = W * H;
165 const input = new Float32Array(3 * plane);
166 let rec: ort.Tensor[] = [0, 1, 2, 3].map(() => new ort.Tensor('float32', new Float32Array([0]), [1, 1, 1, 1]));
167 const dsr = new ort.Tensor('float32', new Float32Array([ds]), [1]);
168 let n = 0;
169 try {
170 for await (const frame of readFrames(decoder.stdout, plane * 3)) {
171 for (let i = 0, j = 0; i < plane; i++, j += 3) {
172 input[i] = frame[j] / 255;
173 input[i + plane] = frame[j + 1] / 255;
174 input[i + 2 * plane] = frame[j + 2] / 255;
175 }
176 const r = await sess.run({
177 src: new ort.Tensor('float32', input, [1, 3, H, W]),
178 r1i: rec[0],
179 r2i: rec[1],
180 r3i: rec[2],
181 r4i: rec[3],
182 downsample_ratio: dsr,
183 });
184 rec = [r.r1o, r.r2o, r.r3o, r.r4o];
185 const fgr = r.fgr.data as Float32Array;
186 const pha = r.pha.data as Float32Array;
187 const rgba = new Uint8ClampedArray(plane * 4);
188 for (let i = 0, j = 0; i < plane; i++, j += 4) {
189 rgba[j] = fgr[i] * 255;
190 rgba[j + 1] = fgr[i + plane] * 255;
191 rgba[j + 2] = fgr[i + 2 * plane] * 255;
192 rgba[j + 3] = pha[i] * 255;
193 }
194 if (!encoder.stdin.write(Buffer.from(rgba.buffer))) await once(encoder.stdin, 'drain');
195 n++;
196 job.progress = Math.min(0.99, n / total);
197 }
198 encoder.stdin.end();
199 const [[dcode], [ecode]] = await Promise.all([decoderDone, encoderDone]);
200 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);
203 return n;
204 } catch (e) {
205 decoder.kill();
206 encoder.kill();
207 await fs.rm(tmp, { force: true });
208 throw e;
209 }
210}