| from torch.utils.data import Dataset | |
| from torchvision.datasets.utils import download_url | |
| from PIL import Image | |
| import torch | |
| import numpy as np | |
| import random | |
| import decord | |
| from decord import VideoReader | |
| import json | |
| import os | |
| from data.utils import pre_caption | |
| decord.bridge.set_bridge("torch") | |
| class ImageNorm(object): | |
| """Apply Normalization to Image Pixels on GPU | |
| """ | |
| def __init__(self, mean, std): | |
| self.mean = torch.tensor(mean).view(1, 3, 1, 1) | |
| self.std = torch.tensor(std).view(1, 3, 1, 1) | |
| def __call__(self, img): | |
| if torch.max(img) > 1 and self.mean.max() <= 1: | |
| img.div_(255.) | |
| return img.sub_(self.mean).div_(self.std) | |
| def load_jsonl(filename): | |
| with open(filename, "r") as f: | |
| return [json.loads(l.strip("\n")) for l in f.readlines()] | |
| class VideoDataset(Dataset): | |
| def __init__(self, video_root, ann_root, num_frm=4, frm_sampling_strategy="rand", max_img_size=384, video_fmt='.mp4'): | |
| ''' | |
| image_root (string): Root directory of video | |
| ann_root (string): directory to store the annotation file | |
| ''' | |
| url = 'https://storage.googleapis.com/sfr-vision-language-research/datasets/msrvtt_test.jsonl' | |
| filename = 'msrvtt_test.jsonl' | |
| download_url(url,ann_root) | |
| self.annotation = load_jsonl(os.path.join(ann_root,filename)) | |
| self.num_frm = num_frm | |
| self.frm_sampling_strategy = frm_sampling_strategy | |
| self.max_img_size = max_img_size | |
| self.video_root = video_root | |
| self.video_fmt = video_fmt | |
| self.img_norm = ImageNorm(mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711)) | |
| self.text = [pre_caption(ann['caption'],40) for ann in self.annotation] | |
| self.txt2video = [i for i in range(len(self.annotation))] | |
| self.video2txt = self.txt2video | |
| def __len__(self): | |
| return len(self.annotation) | |
| def __getitem__(self, index): | |
| ann = self.annotation[index] | |
| video_path = os.path.join(self.video_root, ann['clip_name'] + self.video_fmt) | |
| vid_frm_array = self._load_video_from_path_decord(video_path, height=self.max_img_size, width=self.max_img_size) | |
| video = self.img_norm(vid_frm_array.float()) | |
| return video, ann['clip_name'] | |
| def _load_video_from_path_decord(self, video_path, height=None, width=None, start_time=None, end_time=None, fps=-1): | |
| try: | |
| if not height or not width: | |
| vr = VideoReader(video_path) | |
| else: | |
| vr = VideoReader(video_path, width=width, height=height) | |
| vlen = len(vr) | |
| if start_time or end_time: | |
| assert fps > 0, 'must provide video fps if specifying start and end time.' | |
| start_idx = min(int(start_time * fps), vlen) | |
| end_idx = min(int(end_time * fps), vlen) | |
| else: | |
| start_idx, end_idx = 0, vlen | |
| if self.frm_sampling_strategy == 'uniform': | |
| frame_indices = np.arange(start_idx, end_idx, vlen / self.num_frm, dtype=int) | |
| elif self.frm_sampling_strategy == 'rand': | |
| frame_indices = sorted(random.sample(range(vlen), self.num_frm)) | |
| elif self.frm_sampling_strategy == 'headtail': | |
| frame_indices_head = sorted(random.sample(range(vlen // 2), self.num_frm // 2)) | |
| frame_indices_tail = sorted(random.sample(range(vlen // 2, vlen), self.num_frm // 2)) | |
| frame_indices = frame_indices_head + frame_indices_tail | |
| else: | |
| raise NotImplementedError('Invalid sampling strategy {} '.format(self.frm_sampling_strategy)) | |
| raw_sample_frms = vr.get_batch(frame_indices) | |
| except Exception as e: | |
| return None | |
| raw_sample_frms = raw_sample_frms.permute(0, 3, 1, 2) | |
| return raw_sample_frms | |