Update matte pipeline, server API and video page - proof-of-concept/takes

proof-of-concept / takes

Update matte pipeline, server API and video page

Jeremy Magland <jmagland@flatironinstitute.org> committed 20 days ago commit be26c78d9b6f parent 6805bfc Browse files

7 changed files+365−71

README.mdmodified+4−2View file

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

4646

4747 **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.

4848

49-**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.

49+**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.

5050

51-**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.

51+**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.

52+

53+**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.)

5254

5355 ## Storage layout

5456

server/index.tsmodified+23−5View file

@@ -10,13 +10,13 @@ import os from 'node:os';

1010 import path from 'node:path';

1111 import { fileURLToPath } from 'node:url';

1212 import type { Crop, PublicConfig, TakeCrop, Video } from '../shared/types.ts';

13-import { isFullFrame } from '../shared/types.ts';

13+import { MATTE_RESOLUTIONS, backgroundOf, isFullFrame } from '../shared/types.ts';

1414 import { getExportJob, startExport } from './export.ts';

1515 import { HttpError, checkId, createVideo, getConfig, getVideo, initStore, libraryDir, listVideos, setOrder, trashVideo, updateConfig, updateVideo, videoDir } from './store.ts';

1616 import { appendChunk, createTake, deleteTake, finalizeTake, setThumbnail } from './takes.ts';

1717 import { systemCaptureEvent } from './sysaudio.ts';

1818 import { transcribeTake } from './transcribe.ts';

19-import { requestMattes, withMatteStatus } from './matte.ts';

19+import { cancelMattes, requestMattes, withMatteStatus } from './matte.ts';

2020

2121 function parseArgs(argv: string[]) {

2222 const opts = { dir: path.join(os.homedir(), 'Videos', 'takes'), port: 4455, open: true, browser: false };

@@ -116,7 +116,13 @@ app.patch('/api/videos/:id', json, async (req, res) => {

116116 if (b.layout) {

117117 v.layout = { ...v.layout, ...b.layout };

118118 const bg = v.layout.background;

119- if (bg && (!['none', 'blur', 'color', 'cutout'].includes(bg.mode) || !/^#[0-9a-fA-F]{6}$/.test(bg.color) || !(bg.blur >= 0 && bg.blur <= 100))) {

119+ if (

120+ bg &&

121+ (!['none', 'blur', 'color', 'cutout'].includes(bg.mode) ||

122+ !/^#[0-9a-fA-F]{6}$/.test(bg.color) ||

123+ !(bg.blur >= 0 && bg.blur <= 100) ||

124+ (bg.resolution !== undefined && !MATTE_RESOLUTIONS.includes(bg.resolution)))

125+ ) {

120126 throw new HttpError(400, 'Invalid background settings');

121127 }

122128 }

@@ -135,6 +141,8 @@ app.patch('/api/videos/:id', json, async (req, res) => {

135141 v.timeline = b.timeline.map((s: { id: string; takeId: string; start: number; end: number }) => ({ id: s.id, takeId: s.takeId, start: s.start, end: s.end }));

136142 }

137143 });

144+ // A running background removal job is restarted if the processing resolution changed.

145+ if (b.layout?.background && backgroundOf(v.layout).mode !== 'none') await requestMattes(v.id);

138146 res.json(v);

139147 });

140148

@@ -158,6 +166,7 @@ function parseCrops(raw: unknown, takeIds: Set<string>): Record<string, TakeCrop

158166 }

159167

160168 app.delete('/api/videos/:id', async (req, res) => {

169+ await cancelMattes(p(req, 'id'));

161170 await trashVideo(p(req, 'id'));

162171 res.json({ ok: true });

163172 });

@@ -184,6 +193,7 @@ app.post('/api/videos/:id/takes/:takeId/finalize', async (req, res) => {

184193 });

185194

186195 app.delete('/api/videos/:id/takes/:takeId', async (req, res) => {

196+ await cancelMattes(p(req, 'id'), p(req, 'takeId'));

187197 res.json(withMatteStatus(await deleteTake(p(req, 'id'), p(req, 'takeId'))));

188198 });

189199

@@ -192,10 +202,18 @@ app.post('/api/videos/:id/takes/:takeId/transcribe', async (req, res) => {

192202 res.json(withMatteStatus(await getVideo(p(req, 'id'))));

193203 });

194204

195-// Queue background removal for the video's camera takes (all unprocessed ones, or one take to retry).

205+// Queue background removal for the video's camera takes (all unprocessed ones, or one take to

206+// retry or redo), restarting a running job if the processing resolution has changed.

196207 app.post('/api/videos/:id/matte', json, async (req, res) => {

197208 const takeId = typeof req.body?.takeId === 'string' ? checkId(req.body.takeId) : undefined;

198- await requestMattes(p(req, 'id'), { takeId, retry: !!takeId });

209+ await requestMattes(p(req, 'id'), { takeId });

210+ res.json(withMatteStatus(await getVideo(p(req, 'id'))));

211+});

212+

213+// Cancel queued and running background removal for the video (or one take).

214+app.post('/api/videos/:id/matte/cancel', json, async (req, res) => {

215+ const takeId = typeof req.body?.takeId === 'string' ? checkId(req.body.takeId) : undefined;

216+ await cancelMattes(p(req, 'id'), takeId);

199217 res.json(withMatteStatus(await getVideo(p(req, 'id'))));

200218 });

201219

server/matte-worker.pymodified+11−4View file

@@ -2,7 +2,8 @@

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. Progress is reported on stderr as "takes:progress <frames>" lines.

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

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

67

78 import json

89 import os

@@ -14,10 +15,14 @@ import numpy as np

1415 import onnxruntime as ort

1516

1617

18+def report(line):

19+ print(f'takes:{line}', file=sys.stderr, flush=True)

20+

21+

1722 def ensure_model(url):

1823 path = os.path.expanduser(os.path.join('~', '.cache', 'takes', 'models', os.path.basename(url)))

1924 if not os.path.exists(path):

20- print(f'Downloading background removal model to {path}', file=sys.stderr, flush=True)

25+ report('stage Downloading model')

2126 os.makedirs(os.path.dirname(path), exist_ok=True)

2227 tmp = f'{path}.{os.getpid()}.tmp'

2328 urllib.request.urlretrieve(url, tmp)

@@ -34,6 +39,7 @@ def main():

3439 spec = json.loads(sys.argv[1])

3540 W, H = spec['width'], spec['height']

3641 model = ensure_model(spec['modelUrl'])

42+ report('stage Loading model on the GPU')

3743 if hasattr(ort, 'preload_dlls'):

3844 ort.preload_dlls() # the CUDA and cuDNN libraries installed by pip

3945 opts = ort.SessionOptions()

@@ -42,6 +48,7 @@ def main():

4248 if 'CUDAExecutionProvider' not in sess.get_providers():

4349 sys.exit('onnxruntime could not use the GPU (CUDAExecutionProvider is unavailable); see the messages above')

4450

51+ report('stage Removing background')

4552 decoder = subprocess.Popen(['ffmpeg', *spec['decode']], stdout=subprocess.PIPE)

4653 encoder = subprocess.Popen(['ffmpeg', *spec['encode']], stdin=subprocess.PIPE)

4754 rec = [np.zeros((1, 1, 1, 1), np.float32)] * 4

@@ -64,11 +71,11 @@ def main():

6471 encoder.stdin.write(rgba.tobytes())

6572 n += 1

6673 if n % 10 == 0:

67- print(f'takes:progress {n}', file=sys.stderr, flush=True)

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

6875 encoder.stdin.close()

6976 if decoder.wait() != 0 or encoder.wait() != 0:

7077 sys.exit('ffmpeg failed')

71- print(f'takes:progress {n}', file=sys.stderr, flush=True)

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

7279

7380

7481 if __name__ == '__main__':

server/matte.tsmodified+212−39View file

@@ -3,6 +3,8 @@

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.

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.

68

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

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

@@ -10,10 +12,11 @@ import { createReadStream, createWriteStream } from 'node:fs';

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

1113 import os from 'node:os';

1214 import path from 'node:path';

13-import type { Readable } from 'node:stream';

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

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

1517 import ort from 'onnxruntime-node';

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

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

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

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

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

@@ -24,8 +27,14 @@ const MODEL_PATH = path.join(os.homedir(), '.cache', 'takes', 'models', 'rvm_mob

2427 interface Job {

2528 videoId: string;

2629 takeId: string;

27- progress: number;

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

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

2832 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

2938 }

3039

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

@@ -40,24 +49,66 @@ export const cutoutFile = (videoId: string, takeId: string) => path.join(takeDir

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

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

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

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

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

4453 });

4554 return { ...v, takes };

4655 }

4756

48-// Queue every ready camera take of the video that has not been processed yet (or failed, if retry).

49-export async function requestMattes(videoId: string, opts: { retry?: boolean; takeId?: string } = {}) {

57+const 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

62+

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

64+function 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+}

68+

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.

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

5073 const v = await getVideo(videoId);

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

5175 for (const t of v.takes) {

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

53- if (t.status !== 'ready' || !t.tracks.camera || jobs.has(key(videoId, t.id))) continue;

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

55- if (t.matte?.status === 'failed' && !opts.retry) continue;

56- jobs.set(key(videoId, t.id), { videoId, takeId: t.id, progress: 0, running: false });

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

5791 }

5892 pump();

5993 }

6094

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.

98+export 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+

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

111+

61112 async function pump() {

62113 if (pumping) return;

63114 pumping = true;

@@ -66,22 +117,28 @@ async function pump() {

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

67118 if (!job) break;

68119 job.running = true;

69- const setMatte = (matte: Take['matte']) =>

70- updateVideo(job.videoId, (v) => {

71- const t = v.takes.find((t) => t.id === job.takeId);

72- if (t) t.matte = matte;

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

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

74121 try {

75122 const t0 = Date.now();

76123 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}` : ''}`);

79- await setMatte({ status: 'done', updatedAt: new Date().toISOString() });

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

80128 } catch (e) {

81- console.warn(`matte failed for ${job.videoId}/${job.takeId}:`, (e as Error).message);

82- await setMatte({ status: 'failed', error: (e as Error).message, updatedAt: new Date().toISOString() });

83- } finally {

84- jobs.delete(key(job.videoId, job.takeId));

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+ }

85142 }

86143 }

87144 } finally {

@@ -151,26 +208,31 @@ interface Plan {

151208 total: number;

152209 }

153210

154-async function processTake(job: Job, remote: { host: string; setup: string }): Promise<number> {

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

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

156213 const info = await probe(src);

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

158- // Process at the camera's resolution, capped at 1080 lines.

159- const scale = Math.min(1, 1080 / info.height);

160- const W = even(info.width * scale);

161- const H = even(info.height * scale);

162- // Downsample ratio for the model's coarse pass, following the RVM guidance (0.25 at 1080p, 0.375 at 720p).

163- const ds = Math.min(1, 480 / Math.max(W, H));

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

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

165225

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

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

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

169229 try {

230+ job.abort.signal.throwIfAborted();

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

232+ job.abort.signal.throwIfAborted();

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

172234 await fs.rename(tmp, out);

173- return n;

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

174236 } catch (e) {

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

176238 throw e;

@@ -178,9 +240,22 @@ async function processTake(job: Job, remote: { host: string; setup: string }): P

178240 }

179241

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

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

181244 const sess = await getSession();

245+ job.stage = 'Removing background';

246+ job.progress = 0;

247+ const { signal } = job.abort;

248+ signal.throwIfAborted();

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

183250 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

184259 let stderr = '';

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

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

@@ -194,6 +269,7 @@ async function runLocal(job: Job, { src, out, W, H, ds, total }: Plan): Promise<

194269 let n = 0;

195270 try {

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

272+ signal.throwIfAborted();

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

198274 input[i] = frame[j] / 255;

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

@@ -217,29 +293,37 @@ async function runLocal(job: Job, { src, out, W, H, ds, total }: Plan): Promise<

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

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

219295 }

220- if (!encoder.stdin.write(Buffer.from(rgba.buffer))) await once(encoder.stdin, 'drain');

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

221297 n++;

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

223299 }

300+ signal.throwIfAborted();

224301 encoder.stdin.end();

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

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

227304 return n;

228305 } catch (e) {

229- decoder.kill();

230- encoder.kill();

306+ kill();

231307 throw e;

308+ } finally {

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

232310 }

233311 }

234312

235313 // ---- remote processing over ssh ----

236314 //

237-// A single ssh connection carries the whole job: camera.webm goes to the remote machine on stdin,

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

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

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

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

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

242322 // 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.

243327

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

245329

@@ -249,12 +333,13 @@ function remoteScript(setup: string, worker: string) {

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

250334 } </dev/null

251335 set -e

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

252337 base=~/.cache/takes

253338 venv=$base/venv

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

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

256341 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

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

258343 rm -rf "$venv"

259344 mkdir -p "$base"

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

@@ -267,9 +352,11 @@ trap 'rm -rf "$work"' EXIT

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

268353 ${worker}

269354 TAKES_WORKER

355+echo "takes:upload" >&2

270356 cat > "$work/in.webm"

271357 cd "$work"

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

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

273360 cat out.webm

274361 `;

275362 }

@@ -277,8 +364,24 @@ cat out.webm

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

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

279366 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+ });

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

281371 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+ });

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

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

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

@@ -291,10 +394,21 @@ async function runRemote(job: Job, { src, out, W, H, ds, total }: Plan, host: st

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

292395 partial = parts.pop()!;

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

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

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

399+ n = parseInt(m[2]);

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

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

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

300414 lines.push(line);

@@ -305,13 +419,72 @@ async function runRemote(job: Job, { src, out, W, H, ds, total }: Plan, host: st

305419 try {

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

307421 closed,

308- pipeline(createReadStream(src), ssh.stdin).catch(() => {}), // an early exit shows up in the exit code

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

309423 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+ }),

310429 ]);

311430 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})` : ''}`);

312432 return n;

313433 } catch (e) {

314- ssh.kill();

434+ kill();

315435 throw e;

436+ } finally {

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

316438 }

317439 }

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.

446+async 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+}

shared/types.tsmodified+17−3View file

@@ -15,8 +15,12 @@ export interface Take {

1515 duration: number;

1616 transcript?: { status: 'pending' | 'done' | 'failed'; error?: string; updatedAt: string };

1717 // Background removal of the camera track (cutout.webm, VP9 with alpha). queued/running are live

18- // states reported by the server; only done/failed are persisted.

19- matte?: { status: 'queued' | 'running' | 'done' | 'failed'; progress?: number; error?: string; updatedAt: string };

18+ // states reported by the server; only done/failed/cancelled are persisted. While running, `stage`

19+ // says what is happening and `progress` (absent when unknown) is the fraction of that stage done.

20+ // `height` is the number

21+ // of lines the cutout was processed at, and `capped` says whether that was limited by the

22+ // processing resolution setting rather than by the camera (both absent for older cutouts).

23+ matte?: { status: 'queued' | 'running' | 'done' | 'failed' | 'cancelled'; stage?: string; progress?: number; error?: string; height?: number; capped?: boolean; updatedAt: string };

2024 error?: string;

2125 }

2226

@@ -38,9 +42,19 @@ export interface Background {

3842 mode: BackgroundMode;

3943 color: string; // #rrggbb, for 'color'

4044 blur: number; // Gaussian sigma in pixels at 1080p output, for 'blur'

45+ // Background removal runs on the camera scaled down to at most this many lines. A small

46+ // picture-in-picture needs far fewer lines than the camera records, and the processing time is

47+ // roughly proportional to the number of pixels.

48+ resolution: MatteResolution;

4149 }

4250

43-export const DEFAULT_BACKGROUND: Background = { mode: 'none', color: '#1e2530', blur: 18 };

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

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

53+

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

55+

56+// A cutout processed at fewer lines than the current setting asks for (and fewer than the camera has).

57+export const matteBelowSetting = (m: Take['matte'], bg: Background) => m?.status === 'done' && !!m.capped && (m.height ?? 0) < bg.resolution - 2;

4458

4559 export const backgroundOf = (l: Layout): Background => ({ ...DEFAULT_BACKGROUND, ...l.background });

4660

src/api.tsmodified+1−0View file

@@ -52,6 +52,7 @@ export const api = {

5252 peaks: (id: string, takeId: string) => getOptional<Peaks>(`/media/${id}/takes/${takeId}/peaks.json`),

5353

5454 requestMattes: (id: string, takeId?: string) => request<Video>('POST', `${v(id)}/matte`, takeId ? { takeId } : {}),

55+ cancelMattes: (id: string, takeId?: string) => request<Video>('POST', `${v(id)}/matte/cancel`, takeId ? { takeId } : {}),

5556 uploadThumbnail: (id: string, jpeg: Blob) => request<Video>('POST', `${v(id)}/thumbnail`, jpeg),

5657 startExport: (id: string) => request<ExportJob>('POST', `${v(id)}/export`),

5758 exportStatus: (id: string) => request<ExportJob | null>('GET', `${v(id)}/export`),

src/pages/VideoPage.tsxmodified+97−18View file

@@ -1,6 +1,6 @@

11 import { useCallback, useEffect, useMemo, useRef, useState } from 'react';

22 import type { Background, ExportJob, Layout, OutputSettings, PublicConfig, Segment, Take, TakeCrop, Transcript, Video } from '../../shared/types.ts';

3-import { backgroundOf, exportSignature, systemVolumeOf } from '../../shared/types.ts';

3+import { MATTE_RESOLUTIONS, backgroundOf, cropOf, exportSignature, matteBelowSetting, systemVolumeOf } from '../../shared/types.ts';

44 import { computeGeometry } from '../../shared/layout.ts';

55 import { cutRange, locate, moveSegment, randomId, segDuration, segmentStarts, splitAt, timelineDuration } from '../../shared/timeline.ts';

66 import { navigate } from '../App.tsx';

@@ -328,13 +328,20 @@ export function VideoPage({ videoId }: { videoId: string }) {

328328 const setOutput = (o: Partial<OutputSettings>) => patch({ output: { ...video.output, ...o } }, 200);

329329 const background = backgroundOf(video.layout);

330330 const setBackground = (b: Partial<Background>) => setLayout({ background: { ...background, ...b } });

331- const retryMatte = async (takeId: string) => {

331+ const startMatte = async (takeId: string) => {

332332 try {

333333 refresh(await api.requestMattes(videoId, takeId));

334334 } catch (e) {

335335 setError((e as Error).message);

336336 }

337337 };

338+ const cancelMatte = async (takeId?: string) => {

339+ try {

340+ refresh(await api.cancelMattes(videoId, takeId));

341+ } catch (e) {

342+ setError((e as Error).message);

343+ }

344+ };

338345 const openCrop = (t: Take) => {

339346 // Show the frame under the playhead if it is in this take.

340347 const loc = locate(tl, now());

@@ -357,6 +364,15 @@ export function VideoPage({ videoId }: { videoId: string }) {

357364 patch({ crops }, 200);

358365 };

359366 const cameraTakes = readyTakes.filter((t) => t.tracks.camera);

367+ const matteActive = cameraTakes.some((t) => t.matte?.status === 'queued' || t.matte?.status === 'running');

368+ // The largest camera box in the output and the number of camera lines it shows (more when the

369+ // camera is cropped), as a guide for the processing resolution.

370+ const cameraBox = cameraTakes.reduce<{ w: number; h: number; lines: number } | null>((best, t) => {

371+ const cam = computeGeometry(video.layout, video.output, t.tracks.screen, true).camera;

372+ if (!cam) return best;

373+ const lines = Math.round(cam.box.h / (cropOf(video, t.id, 'camera')?.h ?? 1));

374+ return !best || lines > best.lines ? { w: cam.box.w, h: cam.box.h, lines } : best;

375+ }, null);

360376 // Takes in the timeline whose camera appears in the output but has not had its background removed yet.

361377 const matteBlocking =

362378 background.mode === 'none'

@@ -636,15 +652,42 @@ export function VideoPage({ videoId }: { videoId: string }) {

636652 </label>

637653 )}

638654 {background.mode !== 'none' && (

639- <div className="matte-status">

640- {cameraTakes.map((t) => (

641- <div key={t.id} className="row small">

642- <span>Take {takeNumber(t.id)}:</span>

643- <MatteStatus take={t} onRetry={() => retryMatte(t.id)} />

655+ <>

656+ <div className="field">

657+ <span>Processing resolution</span>

658+ <div className="seg-control">

659+ {MATTE_RESOLUTIONS.map((r) => (

660+ <button key={r} className={background.resolution === r ? 'active' : ''} onClick={() => setBackground({ resolution: r })}>

661+ {r}p

662+ </button>

663+ ))}

644664 </div>

645- ))}

646- <p className="muted small">Background removal runs on this computer (Robust Video Matting) and is done once per take.</p>

647- </div>

665+ </div>

666+ {cameraBox && (

667+ <p className="muted small">

668+ The camera box is {cameraBox.w}×{cameraBox.h} px in the {video.output.height}p output

669+ {cameraBox.lines > cameraBox.h ? `, which shows about ${cameraBox.lines} lines of the cropped camera` : ''}. Processing time grows with the number of pixels, so a resolution a little above that is usually enough.

670+ </p>

671+ )}

672+ <div className="matte-status">

673+ {cameraTakes.map((t) => (

674+ <div key={t.id} className="row small">

675+ <span>Take {takeNumber(t.id)}:</span>

676+ <MatteStatus take={t} background={background} onStart={() => startMatte(t.id)} onCancel={() => cancelMatte(t.id)} />

677+ </div>

678+ ))}

679+ {matteActive && (

680+ <div className="row">

681+ <button className="btn small" onClick={() => cancelMatte()}>

682+ Cancel all

683+ </button>

684+ </div>

685+ )}

686+ <p className="muted small">

687+ Background removal runs on {config?.matteHost ? <code>{config.matteHost}</code> : 'this computer'} (Robust Video Matting) and is done once per take.

688+ </p>

689+ </div>

690+ </>

648691 )}

649692 </>

650693 )}

@@ -729,18 +772,54 @@ export function VideoPage({ videoId }: { videoId: string }) {

729772 );

730773 }

731774

732-function MatteStatus({ take, onRetry }: { take: Take; onRetry: () => void }) {

775+function MatteStatus({ take, background, onStart, onCancel }: { take: Take; background: Background; onStart: () => void; onCancel: () => void }) {

733776 const m = take.matte;

777+ const cancel = (

778+ <button className="btn small" onClick={onCancel}>

779+ Cancel

780+ </button>

781+ );

734782 if (!m) return <span className="muted">waiting</span>;

735- if (m.status === 'done') return <span className="ok-text">done</span>;

736- if (m.status === 'queued') return <span className="muted">queued</span>;

783+ if (m.status === 'done')

784+ return (

785+ <>

786+ <span className="ok-text">done{m.height ? ` (${m.height}p)` : ''}</span>

787+ {matteBelowSetting(m, background) && (

788+ <button className="btn small" onClick={onStart}>

789+ Redo at {background.resolution}p

790+ </button>

791+ )}

792+ </>

793+ );

794+ if (m.status === 'queued')

795+ return (

796+ <>

797+ <span className="muted">queued</span>

798+ {cancel}

799+ </>

800+ );

737801 if (m.status === 'running')

738802 return (

739803 <>

740- <div className="progress mini">

741- <div style={{ width: `${Math.round((m.progress ?? 0) * 100)}%` }} />

742- </div>

743- <span className="muted">{Math.round((m.progress ?? 0) * 100)}%</span>

804+ <span className="muted">{m.stage ?? 'Starting'}</span>

805+ {m.progress !== undefined && (

806+ <>

807+ <div className="progress mini">

808+ <div style={{ width: `${Math.round(m.progress * 100)}%` }} />

809+ </div>

810+ <span className="muted">{Math.round(m.progress * 100)}%</span>

811+ </>

812+ )}

813+ {cancel}

814+ </>

815+ );

816+ if (m.status === 'cancelled')

817+ return (

818+ <>

819+ <span className="muted">cancelled</span>

820+ <button className="btn small" onClick={onStart}>

821+ Start

822+ </button>

744823 </>

745824 );

746825 return (

@@ -748,7 +827,7 @@ function MatteStatus({ take, onRetry }: { take: Take; onRetry: () => void }) {

748827 <span className="error-text" title={m.error}>

749828 failed

750829 </span>

751- <button className="btn small" onClick={onRetry}>

830+ <button className="btn small" onClick={onStart}>

752831 Retry

753832 </button>

754833 </>