馃悕馃悕馃悕
1
2import subprocess
3
4from lib.util import mstreamify, streamify
5
6class VideoWriter:
7 def __init__(self, dims, fps, outfile, logfile):
8 ffmpeg_command = [
9 "ffmpeg",
10 "-f", "rawvideo",
11 "-pix_fmt", "rgb24",
12 "-s", f"{dims[1]}x{dims[0]}",
13 "-r", str(fps),
14 "-i", "-",
15 "-c:v", "libx264",
16 "-pix_fmt", "yuv420p",
17 outfile
18 ]
19
20 self.ffmpeg_log = open(logfile, "a")
21
22 self.ffmpeg_proc = subprocess.Popen(ffmpeg_command,
23 stdin=subprocess.PIPE,
24 stdout=self.ffmpeg_log,
25 stderr=subprocess.STDOUT)
26
27 def close(self):
28 self.ffmpeg_proc.stdin.close()
29 self.ffmpeg_proc.wait()
30 if self.ffmpeg_log:
31 self.ffmpeg_log.close()
32
33 def frame_bytes(self, data):
34 self.ffmpeg_proc.stdin.write(data)
35
36 def mframe(self, mtensor):
37 """ monochrome frame, as torch tensor """
38 self.frame_bytes(mstreamify(mtensor))
39
40 def frame(self, tensor):
41 self.frame_bytes(streamify(tensor))
42
43