add efficient AV decoding #38
Conversation
Apply asyn hotfix for a new bugfix release
Release 4.0.0
…o jbarker/efficient_video
| @@ -31,6 +31,11 @@ dependencies = [ | |||
| "tqdm", | |||
| "webdataset", | |||
| "s3fs", | |||
| "bitstring>=4.2.3", | |||
There was a problem hiding this comment.
Should we add these as extras to be installed like pip install megatron-energon[video]?
Would require the imports in the python files to be allowed to fail though
| video_num_frames: int = 64, | ||
| video_out_frame_size: tuple = (224, 224), |
There was a problem hiding this comment.
Should these be constant and not rather be set when fetching the frames in the user code? I think there should be a dynamic option as well.
E.g. Open-Sora group-batches videos into different sized outputs/frames.
| @@ -0,0 +1,4 @@ | |||
| # Copyright (c) 2024, NVIDIA CORPORATION. | |||
There was a problem hiding this comment.
I'd suggest to move this to a more top-level package. Maybe under megatron.energon.video?
| """ | ||
| extension = re.sub(r".*[.]", "", key) | ||
| # TODO(jbarker): we should add a debug log here | ||
| if extension in "mov mp4 webm mkv".split(): |
There was a problem hiding this comment.
| if extension in "mov mp4 webm mkv".split(): | |
| if extension in ("mov", "mp4", "webm", "mkv"): |
| key: media file extension | ||
| data: raw media bytes | ||
| """ | ||
| extension = re.sub(r".*[.]", "", key) |
There was a problem hiding this comment.
| extension = re.sub(r".*[.]", "", key) | |
| extension = key.rsplit('.', 1)[-1] |
Otherwise, we usually compile regexes before
| out_frame_size=self.video_out_frame_size, | ||
| decode_audio=self.video_decode_audio, | ||
| ) | ||
| elif extension in "flac mp3".split(): |
There was a problem hiding this comment.
| elif extension in "flac mp3".split(): | |
| elif extension in ("flac", "mp3"): |
|
|
||
| DEFAULT_AUDIO_FRAME_SHIFT_MS = 10 # in milliseconds | ||
|
|
||
| class AVDecoder: |
There was a problem hiding this comment.
Maybe we should have an alternative decoder as well, which returns the decoder itself, so the user can decide in user code (=encode_sample) which frames to read?
Like this:
# This function is to be registered as decoder
def read_av_data(key: str, data: bytes):
if key in ("mp3", ...):
return AVData(data)
# This class is now passed to the user's `encode_sample` function (i.e. the raw video
# bytes are essentially passed through). This allows the user to decide on the
# parameters on the fly (e.g. for open-sora).
class AVData:
def __init__(self, raw: bytes):
...
def get_frames(
self,
audio_convert_to_melspec: bool = False,
audio_clip_duration: int = 1,
audio_num_clips: int = -1,
audio_target_rate: int = 16000,
video_decode_audio: bool = False,
video_num_frames: int = 64,
video_out_frame_size: tuple = (224, 224),
) -> AudioVideoData:
...WDYT?
| ) | ||
| return None | ||
|
|
||
| def waveform2melspec(waveform, sample_rate, num_mel_bins, target_length): |
There was a problem hiding this comment.
I feel, the functions below may have their own file, and also reside in the fastseek package?
| class AVDecoder: | ||
| def __init__( | ||
| self, | ||
| audio_convert_to_melspec, |
There was a problem hiding this comment.
Generally, we have all parameters statically typed. Also all class variables are typically typed.
There is still tidying and documentation to be done. I am creating the draft PR to get initial feedback on design and big changes required.