|
|
|
|
|
|
|
import glob |
|
import torch |
|
import torchvision |
|
import matplotlib.pyplot as plt |
|
from torch.utils.data import Dataset |
|
import numpy as np |
|
|
|
|
|
|
|
class MomentsDataset(Dataset): |
|
def __init__(self, videos_folder, num_frames, samples_per_video, frame_size=512) -> None: |
|
super().__init__() |
|
|
|
self.videos_paths = glob.glob(f'{videos_folder}/*mp4') |
|
self.resize = torchvision.transforms.Resize(size=frame_size) |
|
self.center_crop = torchvision.transforms.CenterCrop(size=frame_size) |
|
self.num_samples_per_video = samples_per_video |
|
self.num_frames = num_frames |
|
|
|
def __len__(self): |
|
return len(self.videos_paths) * self.num_samples_per_video |
|
|
|
def __getitem__(self, idx): |
|
video_idx = idx // self.num_samples_per_video |
|
video_path = self.videos_paths[video_idx] |
|
|
|
try: |
|
start_idx = np.random.randint(0, 20) |
|
|
|
unsampled_video_frames, audio_frames, info = torchvision.io.read_video(video_path,output_format="TCHW") |
|
sampled_indices = torch.tensor(np.linspace(start_idx, len(unsampled_video_frames)-1, self.num_frames).astype(int)) |
|
sampled_frames = unsampled_video_frames[sampled_indices] |
|
processed_frames = [] |
|
|
|
for frame in sampled_frames: |
|
resized_cropped_frame = self.center_crop(self.resize(frame)) |
|
processed_frames.append(resized_cropped_frame) |
|
frames = torch.stack(processed_frames, dim=0) |
|
frames = frames.float() / 255.0 |
|
except Exception as e: |
|
print('oops', e) |
|
rand_idx = np.random.randint(0, len(self)) |
|
return self.__getitem__(rand_idx) |
|
|
|
out_dict = {'frames': frames, |
|
'caption': 'none', |
|
'keywords': 'none'} |
|
|
|
return out_dict |
|
|
|
|
|
|
