Skip to content

Instantly share code, notes, and snippets.

@dsuess
Last active January 11, 2019 00:01
Show Gist options
  • Save dsuess/8082a1d41227fd4b0d410509d6eaa595 to your computer and use it in GitHub Desktop.
Save dsuess/8082a1d41227fd4b0d410509d6eaa595 to your computer and use it in GitHub Desktop.
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
@dsuess
Copy link
Author

dsuess commented Jan 11, 2019

Iterator over video(s) based on ffmpeg-python

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment