Class that provides an ImageIO-like interface to a sequence of videos treated as a single video.

multi_video_reader.py

#!/usr/bin/env python3

"""
Class that provides an ImageIO-like interface to a sequence of videos treated as a single video.
"""

import bisect
import copy
import itertools

import imageio

def find_lub_cumsum_index(a, x):
    """
    Find index of least upper bound of `x` in cumulative sum of list `a`.
    """

    a_cumsum = list(itertools.accumulate(a))

    i = bisect.bisect_right(a_cumsum, x)
    if i != len(a):
            return i
    raise ValueError

class MultiVideoReader(object):
    def __init__(self, *file_name_list):
        self._readers = [imageio.get_reader(f) for f in file_name_list]
        
        # Sanity checks to limit support to lists of videos with same format,
        # fps, source_size, and size:
        if len(set([r.format for r in self._readers])) > 1:
            raise RuntimeError('all input readers must have same format')
        meta = [r.get_meta_data() for r in self._readers]
        if len(set([m['plugin'] for m in meta])) > 1:
            raise RuntimeError('all inputs must be read with same plugin')
        if meta[0]['plugin'] != 'ffmpeg':
            raise RuntimeError('only ffmpeg supported')
        if len(set([m['ffmpeg_version'] for m in meta])) > 1:
            raise RuntimeError('all inputs must have same size')
        if len(set([m['fps'] for m in meta])) > 1:
            raise RuntimeError('all inputs must have same fps')
        if len(set([m['source_size'] for m in meta])) > 1:
            raise RuntimeError('all inputs must have same source size')
        if len(set([m['size'] for m in meta])) > 1:
            raise RuntimeError('all inputs must have same size')

        # Set number of frames and duration respectively to sums of frames and
        # durations of specified videos:
        self._format = self._readers[0].format
        self._nframes = [m['nframes'] for m in meta]
        self._nframes_cumsum = list(itertools.accumulate(self._nframes))
        self._duration = [m['duration'] for m in meta]
        self._meta = copy.copy(meta[0])
        self._meta['nframes'] = sum(self._nframes)
        self._meta['duration'] = sum(self._duration)

    @property
    def format(self):
        return self._format

    def get_length(self):
        return self._meta['nframes']

    def get_meta_data(self):
        return self._meta

    def get_data(self, index):
        i = find_lub_cumsum_index(self._nframes, index)
        if i == 0:
            return self._readers[i].get_data(index)
        else:
            return self._readers[i].get_data(index-self._nframes_cumsum[i-1])

    def iter_data(self):
        raise NotImplementedError

    def __iter__(self):
        return self.iter_data()
    
def get_reader(*file_name_list):
    return MultiVideoFormat(file_name_list)