馃悕馃悕馃悕
at main 1.2 kB view raw
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