MahmoodAnaam/LRS2-Validation

import cv2
import torch
import datasets
from torchcodec.decoders import AudioDecoder
from torchcodec.decoders import VideoDecoder

def load_audio(source:str|bytes, start_time:int=0, end_time:int|None=None):
    audio_decoder = AudioDecoder(source)
    if end_time is None:
        end_time = audio_decoder.metadata.duration_seconds_from_header
    waveform = audio_decoder.get_samples_played_in_range(start_time, end_time).data
    return waveform.transpose(1, 0)  # T x 1


def load_video(source:str|bytes, start_time:int=0, end_time:int|None=None):
    video_decoder = VideoDecoder(source, dimension_order="NHWC")
    if end_time is None:
        end_time = video_decoder.metadata.duration_seconds
    vid_rgb = video_decoder.get_frames_played_in_range(start_time, end_time).data
    frames = [cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) for frame in vid_rgb.numpy()]
    vid = torch.from_numpy(np.stack(frames)).unsqueeze(1)
    return vid # T x C x H x W


if __name__=="__main__":

  validation_ds = datasets.load_dataset("MahmoodAnaam/LRS2-Validation", split="validation")
  sample = validation_ds[0]
  audio = load_audio(sample['video']) 
  video = load_video(sample['video']) 
  text = sample['label']

  print(audio.shape) # T x 1
  print(video.shape) # T x C x H x W
  print(text)