video_writer.py

class VideoWriter:
    
    def __init__(self, output_path, overwrite_output=False):
        self.output_path = output_path
        self.active = False
        self._stream = None
        self._shape = None
        self.overwrite_output = overwrite_output
        
    def __enter__(self):
        self.active = True
        return self
        
    def get_stream(self, shape):
        if self._stream is not None:
            if self._shape != shape:
                raise ValueError(f'Shape mismatch {self._shape} != {shape}')
            return self._stream
        if not self.active:
            raise RuntimeError('VideoWriter should be used as a context manager')
            
        height, width = shape
        stream = (
            ffmpeg
            .input('pipe:', format='rawvideo', pix_fmt='rgb24', s=f'{width}x{height}')
            .output(self.output_path, pix_fmt='yuv420p'))
        
        if self.overwrite_output:
            stream = stream.overwrite_output()
        
        self._stream = stream.run_async(pipe_stdin=True)
        self._shape = shape
        return self._stream
        
    def put(self, frames):
        if frames.ndim == 3:
            frames = frames[None]
        _, height, width, colors = frames.shape
        assert colors == 3
        stream = self.get_stream((height, width))
        stream.stdin.write(frames.astype(np.uint8).tobytes())
    
    def __exit__(self, *args):
        self.active = False
        if self._stream is not None:
            self._stream.stdin.close()
            self._stream.wait()
            self._stream = None
            self._shape = None