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, except that the local path also skips the model on some frames
6# with little motion (see matte.ts); here the model is fast, so every distinct frame is processed.
7# A frame identical to the previous one (as when the decoder repeats frames of a camera that
8# delivered fewer than 30 per second) gets the same output without running the model. Stage changes and progress are reported on stderr as
9# "takes:stage <label>" and "takes:progress <frames>" lines.
11import json
12import os
13import subprocess
14import sys
15import urllib.request
17import numpy as np
18import onnxruntime as ort
21def report(line):
22 print(f'takes:{line}', file=sys.stderr, flush=True)
25def ensure_model(url):
26 path = os.path.expanduser(os.path.join('~', '.cache', 'takes', 'models', os.path.basename(url)))
27 if not os.path.exists(path):
28 report('stage Downloading model')
29 os.makedirs(os.path.dirname(path), exist_ok=True)
30 tmp = f'{path}.{os.getpid()}.tmp'
31 urllib.request.urlretrieve(url, tmp)
32 os.rename(tmp, path)
33 return path
36def to_u8(x):
37 # Round half to even and clamp, as a Uint8ClampedArray does in the local path.
38 return np.clip(np.rint(x * 255), 0, 255).astype(np.uint8)
41def main():
42 spec = json.loads(sys.argv[1])
43 W, H = spec['width'], spec['height']
44 model = ensure_model(spec['modelUrl'])
45 report('stage Loading model on the GPU')
46 if hasattr(ort, 'preload_dlls'):
47 ort.preload_dlls() # the CUDA and cuDNN libraries installed by pip
48 opts = ort.SessionOptions()
49 opts.log_severity_level = 3 # errors only; a warning about memcpy nodes is expected for this model
50 sess = ort.InferenceSession(model, opts, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
51 if 'CUDAExecutionProvider' not in sess.get_providers():
52 sys.exit('onnxruntime could not use the GPU (CUDAExecutionProvider is unavailable); see the messages above')
54 report('stage Removing background')
55 decoder = subprocess.Popen(['ffmpeg', *spec['decode']], stdout=subprocess.PIPE)
56 encoder = subprocess.Popen(['ffmpeg', *spec['encode']], stdin=subprocess.PIPE)
57 rec = [np.zeros((1, 1, 1, 1), np.float32)] * 4
58 dsr = np.array([spec['downsampleRatio']], np.float32)
59 rgba = np.empty((H, W, 4), np.uint8)
60 size = W * H * 3
61 n = 0
62 prev = None
63 while True:
64 buf = decoder.stdout.read(size)
65 if len(buf) < size:
66 break
67 if buf == prev:
68 encoder.stdin.write(rgba.tobytes())
69 n += 1
70 if n % 10 == 0:
71 report(f'progress {n}')
72 continue
73 prev = buf
74 frame = np.frombuffer(buf, np.uint8).reshape(H, W, 3)
75 src = (frame.transpose(2, 0, 1)[None] / np.float32(255)).astype(np.float32)
76 fgr, pha, *rec = sess.run(
77 ['fgr', 'pha', 'r1o', 'r2o', 'r3o', 'r4o'],
78 {'src': src, 'r1i': rec[0], 'r2i': rec[1], 'r3i': rec[2], 'r4i': rec[3], 'downsample_ratio': dsr},
79 )
80 rgba[..., :3] = to_u8(fgr[0].transpose(1, 2, 0))
81 rgba[..., 3] = to_u8(pha[0, 0])
82 encoder.stdin.write(rgba.tobytes())
83 n += 1
84 if n % 10 == 0:
85 report(f'progress {n}')
86 encoder.stdin.close()
87 if decoder.wait() != 0 or encoder.wait() != 0:
88 sys.exit('ffmpeg failed')
89 report(f'progress {n}')
92if __name__ == '__main__':
93 main()