307b4e2Run background removal on a remote GPU over sshJeremy Magland 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
969e06cSkip repeated and low-motion frames in background removal; add 180-line processingJeremy Magland 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
be26c78Update matte pipeline, server API and video pageJeremy Magland 9# "takes:stage <label>" and "takes:progress <frames>" lines.
307b4e2Run background removal on a remote GPU over sshJeremy Magland 10
11import json
12import os
13import subprocess
14import sys
15import urllib.request
17import numpy as np
18import onnxruntime as ort
22 print(f'takes:{line}', file=sys.stderr, flush=True)
26 path = os.path.expanduser(os.path.join('~', '.cache', 'takes', 'models', os.path.basename(url)))
27 if not os.path.exists(path):
be26c78Update matte pipeline, server API and video pageJeremy Magland 28 report('stage Downloading model')
307b4e2Run background removal on a remote GPU over sshJeremy Magland 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'])
be26c78Update matte pipeline, server API and video pageJeremy Magland 45 report('stage Loading model on the GPU')
307b4e2Run background removal on a remote GPU over sshJeremy Magland 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')
be26c78Update matte pipeline, server API and video pageJeremy Magland 54 report('stage Removing background')
307b4e2Run background removal on a remote GPU over sshJeremy Magland 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
969e06cSkip repeated and low-motion frames in background removal; add 180-line processingJeremy Magland 62 prev = None
64 buf = decoder.stdout.read(size)
65 if len(buf) < size:
66 break
969e06cSkip repeated and low-motion frames in background removal; add 180-line processingJeremy Magland 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
307b4e2Run background removal on a remote GPU over sshJeremy Magland 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:
87 if decoder.wait() != 0 or encoder.wait() != 0:
88 sys.exit('ffmpeg failed')
307b4e2Run background removal on a remote GPU over sshJeremy Magland 90
92if __name__ == '__main__':
93 main()