dataset.py

import csv
import os
from decord import VideoReader
import torch
import numpy as np

def scale_preserving_aspect_ratio(video, target_short_side):
    _, _, h, w = video.shape
    if h < w:
        new_h = target_short_side
        new_w = int(w / h * target_short_side)
    else:
        new_w = target_short_side
        new_h = int(h / w * target_short_side)
    video = torch.nn.functional.interpolate(video, size=(new_h, new_w), mode='bilinear', align_corners=False)
    return video

COMMON_PREFIXES = [
    "In the video, ",
    "The video shows ",
    "The video features ",
    "The video is ",
    "The video captures ",
    "The image is ",
]

class OpenVid1MDataset(torch.utils.data.Dataset):
    def __init__(self, data_root, csv_path):
        self.data_root = data_root
        csv_path = os.path.join(data_root, csv_path)
        self.data = []
        with open(csv_path, 'r') as f:
            reader = csv.DictReader(f)
            for row in reader:
                self.data.append(row)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        row = self.data[idx]
        video_path = row['video']
        video_path = os.path.join(self.data_root, 'videos', video_path)
        vr = VideoReader(video_path)
        frames = min(len(vr), 49)
        video = vr.get_batch(range(frames)).asnumpy() # shape (T, H, W, C)
        video = np.transpose(video, (3, 0, 1, 2)) # shape (C, T, H, W)
        video = torch.from_numpy(video).float() / 255.0
        video = video * 2 - 1 # scale to [-1, 1]
        
        # scale video to 480p
        video = scale_preserving_aspect_ratio(video, target_short_side=480)
        caption = row['caption']

        # remove common prefixes from caption
        for prefix in COMMON_PREFIXES:
            if caption.startswith(prefix):
                caption = caption[len(prefix):]
                break

        return {
            "video": video,
            "prompt": caption
        }