from torch.utils.data import Dataset import cv2 import h5py import os from tqdm import tqdm import numpy as np import torch from typing import Optional, Tuple, Union # TODO: Try pre-caching video tokens and have dataloader load video tokens instead of frames class VideoHDF5Dataset(Dataset): def __init__( self, video_path: str, transform=None, # postnormalization save_path: Optional[str] = None, train: bool = True, disable_test_split: bool = True, num_frames: int = 4, # context length resize_to: Tuple[int, int] = (64, 64), fps: int = 30, sequence_stride: Optional[int] = None, # default 60//fps fraction_of_dataset: float = 1.0, # fraction of valid starting indices to expose load_chunk_size: int = 1000, # chunk size when reading from HDF5 load_start_index: int = 0, # skip initial frames when reading cached HDF5 preload_ratio: Optional[float] = None, # if set, only load this ratio of cached frames preprocess_read_step: int = 1, # step to subsample raw video during preprocessing preprocess_slice: Optional[Tuple[Union[int, float], Union[int, float]]] = None, # optional slice applied to preprocessed frames; can be (start_idx, end_idx) ints or (start_ratio, end_ratio) floats in [0,1] ) -> None: self.transform = transform self.train = train self.num_frames = num_frames self.fps = fps self.frame_skip = max(1, (sequence_stride if sequence_stride is not None else max(1, 60 // fps))) self.fraction_of_dataset = float(fraction_of_dataset) self.resize_to = resize_to if save_path and os.path.exists(save_path): with h5py.File(save_path, 'r') as h5_file: frames_dset = h5_file['frames'] total = len(frames_dset) n_frames = int(total if preload_ratio is None else max(0, min(total, int(total * preload_ratio)))) self.data = [] for i in tqdm(range(load_start_index, n_frames, load_chunk_size), desc=f"Loading {n_frames} frames"): chunk = frames_dset[i:min(i + load_chunk_size, n_frames)][:] self.data.extend(chunk) self.data = np.array(self.data) else: frames = self._preprocess_video( video_path=video_path, resize_to=resize_to, read_step=preprocess_read_step, slice_spec=preprocess_slice, ) if save_path: print(f"Saving preprocessed frames to {save_path}") with h5py.File(save_path, 'w') as f: f.create_dataset('frames', data=frames, compression='lzf') # Reload into memory to ensure consistent path with h5py.File(save_path, 'r') as h5_file: frames = h5_file['frames'][:] self.data = frames if not disable_test_split: split_idx = int(0.9 * len(self.data)) self.data = self.data[:split_idx] if train else self.data[split_idx:] def _preprocess_video( self, video_path: str, resize_to: Tuple[int, int], read_step: int = 1, slice_spec: Optional[Tuple[Union[int, float], Union[int, float]]] = None, ) -> np.ndarray: print(f"Preprocessing video {video_path}") video = cv2.VideoCapture(video_path) total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) frames = [] step = max(1, int(read_step)) for i in tqdm(range(0, total_frames, step), desc="Processing video frames"): video.set(cv2.CAP_PROP_POS_FRAMES, i) ret, frame = video.read() if not ret: break frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame = cv2.resize(frame, resize_to, interpolation=cv2.INTER_AREA) frames.append(frame) video.release() frames = np.array(frames) if slice_spec is not None and len(frames) > 0: start, end = slice_spec n = len(frames) if isinstance(start, float) or isinstance(end, float): # interpret as ratios s = 0 if start is None else int(n * max(0.0, min(1.0, float(start)))) e = n if end is None else int(n * max(0.0, min(1.0, float(end)))) else: s = 0 if start is None else int(start) e = n if end is None else int(end) s = max(0, min(n, s)) e = max(s, min(n, e)) frames = frames[s:e] return frames def __len__(self) -> int: max_valid_index = int((len(self.data) - (self.num_frames * self.frame_skip)) * self.fraction_of_dataset) return max(0, max_valid_index) def __getitem__(self, index: int): if index >= len(self): raise IndexError(f"Index {index} out of bounds for dataset of length {len(self)}") frame_sequence = self.data[index:index + (self.num_frames * self.frame_skip):self.frame_skip] if len(frame_sequence) != self.num_frames: raise ValueError(f"Expected {self.num_frames} frames, got {len(frame_sequence)} frames") frame_sequence = frame_sequence.astype(np.float32) / 255.0 if self.transform: transformed_frames = [] for frame in frame_sequence: transformed_frame = self.transform(frame) transformed_frames.append(transformed_frame) frame_sequence = torch.stack(transformed_frames, dim=0) else: frame_sequence = torch.from_numpy(frame_sequence).permute(0, 3, 1, 2) return frame_sequence, 0 def __del__(self): if hasattr(self, 'h5_file'): self.h5_file.close() # TODO: add more datasets class PongDataset(VideoHDF5Dataset): def __init__(self, video_path, transform=None, save_path=None, train=True, num_frames=1, resolution=(64, 64), fps=30, preload_ratio=1): super().__init__( video_path=video_path, transform=transform, save_path=save_path, train=train, num_frames=num_frames, resize_to=resolution, fps=fps, preload_ratio=preload_ratio, load_chunk_size=1000, load_start_index=0, preprocess_read_step=10, # keep every 10th frame from raw preprocess_slice=None, ) class PolePositionDataset(VideoHDF5Dataset): def __init__(self, video_path, transform=None, save_path=None, train=True, num_frames=4, resolution=(64, 64), fps=30, preload_ratio=1): super().__init__( video_path=video_path, transform=transform, save_path=save_path, train=train, num_frames=num_frames, resize_to=resolution, fps=fps, preload_ratio=preload_ratio, sequence_stride=None, load_chunk_size=1000, load_start_index=0, preprocess_read_step=1, preprocess_slice=(1/50, 1/4), ) class SonicDataset(VideoHDF5Dataset): def __init__(self, video_path, transform=None, save_path=None, train=True, num_frames=4, resolution=(128, 128), fps=15, preload_ratio=1): super().__init__( video_path=video_path, transform=transform, save_path=save_path, train=train, num_frames=num_frames, resize_to=resolution, fps=fps, preload_ratio=preload_ratio, sequence_stride=None, load_chunk_size=1000, load_start_index=100, preprocess_read_step=1, preprocess_slice=None, ) class PicoDoomDataset(VideoHDF5Dataset): def __init__(self, video_path, transform=None, save_path=None, train=True, num_frames=4, resolution=(128, 128), fps=30, preload_ratio=0.3): super().__init__( video_path=video_path, transform=transform, save_path=save_path, train=train, num_frames=num_frames, resize_to=resolution, fps=fps, preload_ratio=preload_ratio, sequence_stride=None, load_chunk_size=1000, load_start_index=300, preprocess_read_step=1, preprocess_slice=None, ) class ZeldaDataset(VideoHDF5Dataset): def __init__(self, video_path, transform=None, save_path=None, train=True, num_frames=4, resolution=(128, 128), fps=15, preload_ratio=0.2): super().__init__( video_path=video_path, transform=transform, save_path=save_path, train=train, num_frames=num_frames, resize_to=resolution, fps=fps, preload_ratio=preload_ratio, sequence_stride=None, load_chunk_size=1000, load_start_index=1000, preprocess_read_step=1, preprocess_slice=None, )