From 009fb15d274d36b6e6755a890449975e863ce004 Mon Sep 17 00:00:00 2001 From: "junjie.jiang" Date: Mon, 19 Sep 2022 16:40:30 +0800 Subject: [PATCH] support batch Signed-off-by: junjie.jiang --- __init__.py | 4 ++-- audio_decoder_ffmpeg.py | 31 +++++++++++++++++++++++-------- 2 files changed, 25 insertions(+), 10 deletions(-) diff --git a/__init__.py b/__init__.py index f1d9302..0cf02a5 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,5 @@ from .audio_decoder_ffmpeg import AudioDecoderFFmpeg -def ffmpeg(): - return AudioDecoderFFmpeg() +def ffmpeg(batch_size=-1): + return AudioDecoderFFmpeg(batch_size) diff --git a/audio_decoder_ffmpeg.py b/audio_decoder_ffmpeg.py index ca9b582..d8237bd 100644 --- a/audio_decoder_ffmpeg.py +++ b/audio_decoder_ffmpeg.py @@ -11,15 +11,30 @@ class AudioDecoderFFmpeg(PyOperator): """ """ - def __init__(self) -> None: + def __init__(self, batch_size=-1) -> None: super().__init__() + self._batch_size = batch_size def __call__(self, audio_path: str): + frames = [] in_container = av.open(audio_path) - stream = in_container.streams.get(audio=0)[0] - for frame in in_container.decode(stream): - timestamp = int(frame.time * 1000) - sample_rate = frame.sample_rate - layout = frame.layout.name - ndarray = frame.to_ndarray() - yield AudioFrame(ndarray, sample_rate, timestamp, layout) + stream = in_container.streams.get(audio=0)[0] + if self._batch_size <= 0: + for frame in in_container.decode(stream): + timestamp = int(frame.time * 1000) + sample_rate = frame.sample_rate + layout = frame.layout.name + ndarray = frame.to_ndarray() + yield AudioFrame(ndarray, sample_rate, timestamp, layout) + else: + for frame in in_container.decode(stream): + timestamp = int(frame.time * 1000) + sample_rate = frame.sample_rate + layout = frame.layout.name + ndarray = frame.to_ndarray() + frames.append(AudioFrame(ndarray, sample_rate, timestamp, layout)) + if len(frames) == self._batch_size: + yield frames + frames = [] + if len(frames) != 0: + yield frames