Spaces:
Running on Zero
Running on Zero
File size: 9,152 Bytes
dcc4594 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 | 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,
)
|