video_reader.py

class VideoReader:
    
    def __init__(self, *paths, sampling_rate=None, batch_size=None, drop_last_batch=False):
        self.paths = list(paths)
        self.sampling_rate = sampling_rate
        self.batch_size = batch_size
        self.drop_last_batch = drop_last_batch
    
    @staticmethod
    def _get_video_properties(path):
        probe = ffmpeg.probe(path)
        try:
            video_stream = next(stream for stream in probe['streams'] if stream['codec_type'] == 'video')
        except StopIteration:
            raise ValueError(f'No video stream found in {path}')
        return video_stream
    
    def get_dimensions(self):
        props = [self._get_video_properties(p) for p in self.paths]
        widths = set(v['width'] for v in props)
        heights = set(v['height'] for v in props)
        if len(widths) > 1 or len(heights) > 1:
            raise ValueError(f'Cannot match widths={widths} and heights={heights}')
        return tuple(heights)[0], tuple(widths)[0]
        
    def get_stream(self):
        r = 1 if self.sampling_rate is None else self.sampling_rate
        stream = ffmpeg.concat(*[ffmpeg.input(path, r=r) for path in self.paths])
        
        if self.sampling_rate is not None:
            stream = stream.filter('select', f'not(mod(n,{r}))')
        return stream
        
    def __iter__(self):
        height, width = self.get_dimensions()
        stream = self.get_stream()
        read_process = (
            stream
            .output('pipe:', format='rawvideo', pix_fmt='rgb24', r=1)
            .run_async(pipe_stdout=True))
        
        batch = 1 if self.batch_size is None else self.batch_size
        while True:
            in_bytes = read_process.stdout.read(batch * height * width * 3)
            if not in_bytes:
                break
            frames = np.frombuffer(in_bytes, dtype=np.uint8)
            frames = frames.reshape((-1, height, width, 3))
            
            if self.drop_last_batch and len(frames) < batch:
                break
            yield frames if self.batch_size is not None else frames[0]
        
        read_process.wait()
        return

Iterator over video(s) based on ffmpeg-python