Spaces:
Running on Zero
Running on Zero
File size: 1,510 Bytes
0122a25 | 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 | """Video Dataset Wrapper."""
from __future__ import annotations
from torch.utils.data import Dataset
from .const import CommonKeys as K
from .datasets.base import VideoDataset
from .typing import DictData
class SequentialVideoDataset(Dataset[list[DictData]]):
"""VideoDataset wrapper to have sequential inputs."""
def __init__(self, dataset: VideoDataset, num_frames: int = 2):
"""Init."""
self.dataset = dataset
# Video settings
self.video_mapping = dataset.video_mapping
self.num_frames = num_frames
self.has_reference = True
def __len__(self) -> int:
"""Get length."""
return len(self.dataset)
def __getitem__(self, idx: int) -> list[DictData]:
"""Get item."""
cur_sample = self.dataset[idx]
indices_in_video = self.video_mapping["video_to_indices"][
cur_sample[K.sequence_names]
]
frame_ids = self.video_mapping["video_to_frame_ids"][
cur_sample[K.sequence_names]
]
cur_frame_id = frame_ids[indices_in_video.index(idx)]
samples = []
for i in range(self.num_frames - 1, 0, -1):
past_frame_id = cur_frame_id - i
if past_frame_id >= 0:
samples.append(self.dataset[indices_in_video[past_frame_id]])
# else:
# samples.append(self.dataset[idx])
# Append current sample as the last frame
samples.append(cur_sample)
return samples
|