Helper class for recording frames into video file

videowriter.py

import cv2

VIDEO_WRITER_FOURCC = "H264"

class Writer:
    """
    Usage:
    
        writer = Writer("myvideofile.mp4")
        for frame in <source>:
            writer.write_frame(frame)
        writer.close()
    
    Alternative, as context manager:
    
        with Writer("myvideofile.mp4") as writer:
            for frame in <source>:
                writer.write_frame(frame)
        
    """
    
    def __init__(self, output):
        self.output = output
        self.writer = None
        self.fourcc = cv2.VideoWriter_fourcc(*VIDEO_WRITER_FOURCC)

    def __enter__(self):
        return self

    def write_frame(self, frame):
        if self.writer is None:
            (H, W) = frame.shape[:2]
            self.writer = cv2.VideoWriter(self.output, self.fourcc, 30, (W, H), True)
        self.writer.write(frame)

    def __exit__(self, *args, **kwargs):
        self.close()

    def close(self):
        if self.writer is not None:
            self.writer.release()
            self.writer = None