Spaces:
Running on Zero
Running on Zero
| 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, | |
| ) | |