datasets/video_base.py

from pathlib import Path
from typing import Any, Dict, List, Tuple
import random
import threading
import pandas as pd
import numpy as np
from tqdm import tqdm
import torch
from omegaconf import DictConfig
from torch.utils.data import Dataset
from torchvision.transforms import v2 as transforms


# Must import after torch because this can sometimes lead to a nasty segmentation fault, or stack smashing error
# Very few bug reports but it happens. Look in decord Github issues for more relevant information.
import decord  # isort:skip

decord.bridge.set_bridge("torch")


class VideoDataset(Dataset):
    def __init__(self, cfg: DictConfig, split: str = "training") -> None:
        super().__init__()
        self.cfg = cfg
        self.debug = cfg.debug
        self.split = split
        self.data_root = Path(cfg.data_root)
        self.metadata_path = Path(cfg.metadata_path)
        self.auto_download = cfg.auto_download
        self.force_download = cfg.force_download
        self.test_percentage = cfg.test_percentage
        self.id_token = cfg.id_token or ""
        self.height = cfg.height
        self.width = cfg.width
        self.n_frames = cfg.n_frames
        self.fps = cfg.fps
        self.trim_mode = cfg.trim_mode
        self.pad_mode = cfg.pad_mode
        self.filtering = cfg.filtering
        self.load_video_latent = cfg.load_video_latent
        self.load_prompt_embed = cfg.load_prompt_embed
        self.augmentation = cfg.augmentation
        self.image_to_video = cfg.image_to_video
        self.max_text_tokens = cfg.max_text_tokens

        # trigger auto-download if not already downloaded
        trigger_download = False
        if not self.data_root.is_dir():
            print(f"Dataset root folder {self.data_root} does not exist.")
            if not self.auto_download:
                raise ValueError(
                    f"Attempting to automatically download the dataset since dataset root folder {self.data_root} does not exist. "
                    "If this is the intended behavior, append `dataset.auto_download=True` in your command to pass this check."
                )
            trigger_download = True
        if self.force_download:
            trigger_download = True
        if trigger_download:
            # if threading.current_thread() is not threading.main_thread():
            if torch.distributed.is_initialized():
                raise ValueError(
                    "Download must be called from the main thread with single-process training. Did you call this inside a multi-worker dataloader?"
                )
            print(f"Attempting to download dataset to {self.data_root}...")
            self.download()

        self.records = self._load_records()  # a list of dictionaries
        self.augment_transforms = self._build_video_transforms(augment=True)
        self.no_augment_transforms = self._build_video_transforms(augment=False)
        self.img_normalize = transforms.Normalize(
            mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True
        )

        if self.trim_mode not in ["speedup", "random_cut"]:
            raise ValueError(
                f"Invalid trim_mode: {self.trim_mode}. Must be one of ['speedup', 'random_cut']."
            )
        if self.pad_mode not in ["slowdown", "pad_last", "discard"]:
            raise ValueError(
                f"Invalid pad_mode: {self.pad_mode}. Must be one of ['slowdown', 'pad_last', 'discard']."
            )

    def _build_video_transforms(self, augment: bool = True):
        trans = []
        if augment and self.augmentation.random_flip is not None:
            trans.append(transforms.RandomHorizontalFlip(self.augmentation.random_flip))

        aspect_ratio = self.width / self.height
        aspect_ratio = [aspect_ratio, aspect_ratio]
        if augment and self.augmentation.ratio is not None:
            aspect_ratio[0] *= self.augmentation.ratio[0]
            aspect_ratio[1] *= self.augmentation.ratio[1]

        scale = [1.0, 1.0]
        if augment and self.augmentation.scale is not None:
            scale[0] *= self.augmentation.scale[0]
            scale[1] *= self.augmentation.scale[1]

        trans.append(
            transforms.RandomResizedCrop(
                size=(self.height, self.width),
                scale=scale,
                ratio=aspect_ratio,
                interpolation=transforms.InterpolationMode.BICUBIC,
            ),
        )
        return transforms.Compose(trans)

    def preprocess_record(self, record: Dict[str, Any]) -> Dict[str, Any]:
        # a hook to modify the original record on the fly
        return record

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        record = self.records[idx]

        # Load video data - either raw or preprocessed latents
        videos = self._load_video(record)
        # images = videos[:1].clone() if self.image_to_video else None
        image_latents, video_latents = None, None
        video_metadata = {
            "num_frames": videos.shape[0],
            "height": videos.shape[2],
            "width": videos.shape[3],
        }

        if self.load_video_latent:
            image_latents, video_latents = self._load_video_latent(record)
            # This is hardcoded for now.
            # The VAE's temporal compression ratio is 4.
            # The VAE's spatial compression ratio is 8.
            latent_num_frames = video_latents.size(1)
            if latent_num_frames % 2 == 0:
                n_frames = latent_num_frames * 4
            else:
                n_frames = (latent_num_frames - 1) * 4 + 1

            height = video_latents.size(2) * 8
            width = video_latents.size(3) * 8

            assert video_metadata["num_frames"] == n_frames, "num_frames changed"
            assert video_metadata["height"] == height, "height changed"
            assert video_metadata["width"] == width, "width changed"

        # Load prompt data - either raw or preprocessed embeddings
        caption = ""
        if "caption" in record:
            caption = record["caption"]
        elif "gemini_caption" in record:
            caption = record["gemini_caption"]
        elif "original_caption" in record:
            caption = record["original_caption"]
        video_metadata["has_caption"] = caption != ""
        prompts = self.id_token + caption
        prompt_embeds = None
        prompt_embed_len = None
        if self.load_prompt_embed:
            prompt_embeds, prompt_embed_len = self._load_prompt_embed(record)

        has_bbox, bbox_render = self._render_bbox(record)

        output = {
            "videos": videos,
            "video_metadata": video_metadata,
            "bbox_render": bbox_render,
            "has_bbox": has_bbox,
        }

        if prompts is not None:
            output["prompts"] = prompts
        # if images is not None:
        #     output["images"] = images
        if prompt_embeds is not None:
            output["prompt_embeds"] = prompt_embeds
            output["prompt_embed_len"] = prompt_embed_len
        if image_latents is not None:
            output["image_latents"] = image_latents
        if video_latents is not None:
            output["video_latents"] = video_latents

        return output

    def _n_frames_in_src(self, src_fps):
        """
        Given the fps of the source video, return the number of frames in it we shall
        use in order to generate a target video of self.n_frames frames at self.fps.

        Note the definition of fps of the source video is described in README.md as,
        for a real-world task that requires 1 second to finish, how many frames does it
        take this source video to capture? This is usually just the fps of the source
        video, but if the source video is already a slow motion video, this may be
        different.
        """
        return round(self.n_frames / self.fps * src_fps)

    def _temporal_sample(self, n_frames: int, fps: int) -> torch.Tensor:
        """
        Given number of frames and fps, return a sequence of frame indices to downsample / upsample the video temporally.
        This shall consider self.n_frames and fps.
        """

        # target_len is the number of frames in the source video that we shall use to generate a target video of self.n_frames frames at self.fps
        target_len = self._n_frames_in_src(fps)

        if n_frames < target_len:
            if self.pad_mode == "pad_last":
                indices = np.linspace(0, target_len - 1, self.n_frames)
                indices = np.clip(indices, 0, n_frames - 1)
            elif self.pad_mode == "slowdown":
                indices = np.linspace(0, n_frames - 1, self.n_frames)
            elif self.pad_mode == "discard":
                raise ValueError(
                    "pad_mode is set to 'discard', but this short video is not filtered out."
                )
            else:
                raise ValueError(f"Invalid pad_mode: {self.pad_mode}")
        elif n_frames > target_len:
            if self.trim_mode == "random_cut":
                start = np.random.randint(0, n_frames - target_len)
                indices = start + np.linspace(0, target_len - 1, self.n_frames)
            elif self.trim_mode == "speedup":
                indices = np.linspace(0, n_frames - 1, self.n_frames)
            elif self.trim_mode == "discard":
                raise ValueError(
                    "trim_mode is set to 'discard', but this long video is not filtered out."
                )
            else:
                raise ValueError(f"Invalid trim_mode: {self.trim_mode}")
        else:
            indices = np.linspace(0, n_frames - 1, self.n_frames)

        indices = np.round(indices).astype(int)
        return indices

    def _load_video(self, record: Dict[str, Any]) -> torch.Tensor:
        """
        Given a record, return a tensor of shape (n_frames, 3, H, W)
        """

        video_path = self.data_root / record["video_path"]
        video_reader = decord.VideoReader(uri=video_path.as_posix())
        n_frames = len(video_reader)
        start = record.get("trim_start", 0)
        end = record.get("trim_end", n_frames)
        indices = self._temporal_sample(end - start, record["fps"])
        indices = list(start + indices)
        frames = video_reader.get_batch(indices)

        # do some padding
        if len(frames) != self.n_frames:
            raise ValueError(
                f"Expected {len(frames)=} to be equal to {self.n_frames=}."
            )

        # crop if specified in the record
        if "crop_top" in record and "crop_bottom" in record:
            frames = frames[:, record["crop_top"] : record["crop_bottom"]]
        if "crop_left" in record and "crop_right" in record:
            frames = frames[:, :, record["crop_left"] : record["crop_right"]]

        frames = frames.float().permute(0, 3, 1, 2).contiguous() / 255.0

        if "has_bbox" in record and record["has_bbox"]:
            frames = self.no_augment_transforms(frames)
        else:
            frames = self.augment_transforms(frames)
        frames = self.img_normalize(frames)

        return frames

    def _render_bbox(self, record: Dict[str, Any]) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Given a record, return a tensor of shape (H, W)
        """

        # first frame and last frame forms 2 channels
        bbox_render = torch.zeros(2, record["height"], record["width"])
        has_bbox = torch.zeros(2, dtype=torch.bool)
        # if "first_frame_has_bbox" in record and record["first_frame_has_bbox"]:
        #     has_bbox[0] = True
        #     bbox_top = int(record["first_frame_bbox_top"])
        #     bbox_bottom = int(record["first_frame_bbox_bottom"])
        #     bbox_left = int(record["first_frame_bbox_left"])
        #     bbox_right = int(record["first_frame_bbox_right"])
        #     bbox_render[0, bbox_top:bbox_bottom, bbox_left:bbox_right] = 1
        # if "last_frame_has_bbox" in record and record["last_frame_has_bbox"]:
        #     has_bbox[-1] = True
        #     bbox_top = int(record["last_frame_bbox_top"])
        #     bbox_bottom = int(record["last_frame_bbox_bottom"])
        #     bbox_left = int(record["last_frame_bbox_left"])
        #     bbox_right = int(record["last_frame_bbox_right"])
        #     bbox_render[-1, bbox_top:bbox_bottom, bbox_left:bbox_right] = 1
        bbox_render = self.no_augment_transforms(bbox_render)
        return has_bbox, bbox_render

    def _load_records(self) -> Tuple[List[str], List[str]]:
        """
        Given the metadata file, loads the records as a list.
        Each record is a dictionary containing a datapoint's video path / caption etc.
        Require these entries: "video_path", "caption", "height", "width", "n_frames", "fps"
        Optional entry: "split" - if present, will be used instead of test_percentage
        """

        records = pd.read_csv(self.data_root / self.metadata_path, na_filter=False)
        records = records.to_dict("records")
        len_pre_filter = len(records)
        if not self.filtering.disable:
            records = [record for record in records if self._filter_record(record)]
        len_post_filter = len(records)

        print(
            f"{self.data_root / self.metadata_path}: filtered {len_pre_filter - len_post_filter} records from {len_pre_filter} to {len_post_filter}, rataining rate: {len_post_filter / len_pre_filter}"
        )

        if self.cfg.check_video_path and not self.debug:
            print("Checking records such that all video_path are valid...")
            print(
                "This could take a while. To skip, append `dataset.check_video_path=False` to your command."
            )
            for r in tqdm(records, desc="Checking video paths"):
                self._check_record(r)
            print("Done checking records")

        # Handle split selection
        if self.split != "all":
            if "split" in records[0]:
                # Use split field from records
                records = [r for r in records if r["split"] == self.split]
                if not records:
                    raise ValueError(f"No records found for split '{self.split}'")
            else:
                # Use test_percentage
                if self.split == "training":
                    records = records[: -int(len(records) * self.test_percentage)]
                else:  # validation/test
                    records = records[-int(len(records) * self.test_percentage) :]

        random.Random(0).shuffle(records)

        records = [self.preprocess_record(record) for record in records]

        return records

    def _filter_record(self, x: Dict[str, Any]) -> bool:
        """
        x is a record dictionary containing a datapoint's video path / caption etc.
        Returns True if the record should be kept, False otherwise.
        """
        h, w, fps = x["height"], x["width"], x["fps"]

        # if record specified a crop, use that
        if "crop_left" in x and "crop_right" in x:
            w = x["crop_right"] - x["crop_left"]
        if "crop_top" in x and "crop_bottom" in x:
            h = x["crop_bottom"] - x["crop_top"]
        if "trim_start" in x and "trim_end" in x:
            n_frames = x["trim_end"] - x["trim_start"]
        elif "n_frames" in x:
            n_frames = x["n_frames"]
        else:
            raise ValueError(
                "Record missing required key 'n_frames', if trim not specified"
            )

        h_range = self.filtering.height
        if h_range is not None and h < h_range[0] or h > h_range[1]:
            return False
        w_range = self.filtering.width
        if w_range is not None and w < w_range[0] or w > w_range[1]:
            return False
        f_range = self.filtering.n_frames
        if f_range is not None and n_frames < f_range[0] or n_frames > f_range[1]:
            return False
        fps_range = self.filtering.fps
        if fps_range is not None and fps < fps_range[0] or fps > fps_range[1]:
            return False
        if n_frames < self._n_frames_in_src(fps) and self.pad_mode == "discard":
            return False

        # then filter using stable_background, stable_brightness,
        # note that some datasets may not have these keys
        if "stable_background" in x and not x["stable_background"]:
            return False
        if "stable_brightness" in x and not x["stable_brightness"]:
            return False

        return True

    def _check_record(self, x: Dict[str, Any]) -> bool:
        """
        x is a record dictionary containing a datapoint's video path / caption etc.
        raise an error if the record is not valid. e.g.
        """
        video_path = self.data_root / x["video_path"]
        if not video_path.is_file():
            msg = f"Expected `{video_path=}` to be a valid file but found it to be invalid."
            if self.debug:
                print(msg)
            else:
                raise ValueError(msg)

    def _load_video_latent(
        self, record: Dict[str, Any]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if "video_latent_path" not in record:
            raise ValueError("Record missing required key 'video_latent_path'")
        video_latent_path = self.data_root / record["video_latent_path"]

        image_latent = None
        if self.image_to_video:
            if "image_latent_path" not in record:
                raise ValueError("Record missing required key 'image_latent_path'")
            image_latent_path = self.data_root / record["image_latent_path"]
            image_latent = torch.load(
                image_latent_path, map_location="cpu", weights_only=True
            )
        video_latent = torch.load(
            video_latent_path, map_location="cpu", weights_only=True
        )

        return image_latent, video_latent

    def _load_prompt_embed(self, record: Dict[str, Any]) -> torch.Tensor:
        # if self.debug:
        #     return torch.zeros(self.max_text_tokens, 4096), self.max_text_tokens

        if "prompt_embed_path" not in record:
            raise ValueError("Record missing required key 'prompt_embed_path'")
        prompt_embed_path = self.data_root / record["prompt_embed_path"]
        prompt_embed = torch.load(
            prompt_embed_path, map_location="cpu", weights_only=True
        )

        prompt_embed_len = prompt_embed.size(0)
        if prompt_embed_len < self.max_text_tokens:
            # Pad with zeros to max_text_tokens
            padding = torch.zeros(
                self.max_text_tokens - prompt_embed.size(0),
                prompt_embed.size(1),
                dtype=prompt_embed.dtype,
                device=prompt_embed.device,
            )
            prompt_embed = torch.cat([prompt_embed, padding], dim=0)

        return prompt_embed, prompt_embed_len

    def download(self):
        """
        Automatically download the dataset to self.data_root. Optional.
        """
        raise NotImplementedError(
            "Automatic download not implemented for this dataset."
        )