| import csv |
| import os |
|
|
| import numpy as np |
| import torch |
| import torchvision |
| import torchvision.transforms as transforms |
| from torchvision.datasets.folder import IMG_EXTENSIONS, pil_loader |
|
|
| from . import video_transforms |
| from .utils import center_crop_arr |
|
|
|
|
| def get_transforms_video(resolution=256): |
| transform_video = transforms.Compose( |
| [ |
| video_transforms.ToTensorVideo(), |
| video_transforms.RandomHorizontalFlipVideo(), |
| video_transforms.UCFCenterCropVideo(resolution), |
| transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), |
| ] |
| ) |
| return transform_video |
|
|
|
|
| def get_transforms_image(image_size=256): |
| transform = transforms.Compose( |
| [ |
| transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, image_size)), |
| transforms.RandomHorizontalFlip(), |
| transforms.ToTensor(), |
| transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), |
| ] |
| ) |
| return transform |
|
|
|
|
| class DatasetFromCSV(torch.utils.data.Dataset): |
| """load video according to the csv file. |
| |
| Args: |
| target_video_len (int): the number of video frames will be load. |
| align_transform (callable): Align different videos in a specified size. |
| temporal_sample (callable): Sample the target length of a video. |
| """ |
|
|
| def __init__( |
| self, |
| csv_path, |
| num_frames=16, |
| frame_interval=1, |
| transform=None, |
| root=None, |
| ): |
| self.csv_path = csv_path |
| with open(csv_path, "r") as f: |
| reader = csv.reader(f) |
| self.samples = list(reader) |
|
|
| ext = self.samples[0][0].split(".")[-1] |
| if ext.lower() in ("mp4", "avi", "mov", "mkv"): |
| self.is_video = True |
| else: |
| assert f".{ext.lower()}" in IMG_EXTENSIONS, f"Unsupported file format: {ext}" |
| self.is_video = False |
|
|
| self.transform = transform |
|
|
| self.num_frames = num_frames |
| self.frame_interval = frame_interval |
| self.temporal_sample = video_transforms.TemporalRandomCrop(num_frames * frame_interval) |
| self.root = root |
|
|
| def getitem(self, index): |
| sample = self.samples[index] |
| path = sample[0] |
| if self.root: |
| path = os.path.join(self.root, path) |
| text = sample[1] |
|
|
| if self.is_video: |
| vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit="sec", output_format="TCHW") |
| total_frames = len(vframes) |
|
|
| |
| start_frame_ind, end_frame_ind = self.temporal_sample(total_frames) |
| assert ( |
| end_frame_ind - start_frame_ind >= self.num_frames |
| ), f"{path} with index {index} has not enough frames." |
| frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int) |
|
|
| video = vframes[frame_indice] |
| video = self.transform(video) |
| else: |
| image = pil_loader(path) |
| image = self.transform(image) |
| video = image.unsqueeze(0).repeat(self.num_frames, 1, 1, 1) |
|
|
| |
| video = video.permute(1, 0, 2, 3) |
|
|
| return {"video": video, "text": text} |
|
|
| def __getitem__(self, index): |
| for _ in range(10): |
| try: |
| return self.getitem(index) |
| except Exception as e: |
| print(e) |
| index = np.random.randint(len(self)) |
| raise RuntimeError("Too many bad data.") |
|
|
| def __len__(self): |
| return len(self.samples) |
|
|
