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