frames_dataset.py

import os
from skimage import io, img_as_float32
from skimage.color import gray2rgb
from sklearn.model_selection import train_test_split
from imageio import mimread

import numpy as np
from torch.utils.data import Dataset, dataset
import pandas as pd
from augmentation import AllAugmentationTransform
import glob


def read_video(root_dir, frame_shape=None):
    """
    Read video which can be:
      - an image of concatenated frames
      - '.mp4' and'.gif'
      - folder with videos
    """
    # Check if frame_shape is an integer
    if frame_shape is not None and not all(isinstance(i, int) for i in frame_shape):
        raise ValueError("frame_shape must be a tuple of integers")

    import yaml

    # Load the configuration file
    with open("config.yaml", "r") as f:
        config = yaml.safe_load(f)

    # Get the root directory from the configuration file
    root_dir = config["root_dir"]

    if os.path.isdir(root_dir):
        frames = sorted(os.listdir(root_dir))
        num_frames = len(frames)
        video_array = np.asarray(
            [
                img_as_float32(io.imread(os.path.join(root_dir, frames[idx])))
                for idx in range(num_frames)
            ]
        )
    elif root_dir.lower().endswith(".png") or root_dir.lower().endswith(".jpg"):
        image = io.imread(root_dir)

        if len(image.shape) == 2 or image.shape[2] == 1:
            image = gray2rgb(image)

        if image.shape[2] == 4:
            image = image[..., :3]

        image = img_as_float32(image)

        video_array = np.moveaxis(image, 1, 0)

        video_array = np.asarray(video_array).reshape((-1,) + frame_shape)
        video_array = np.moveaxis(video_array, 1, 2)
    elif (
        root_dir.lower().endswith(".gif")
        or root_dir.lower().endswith(".mp4")
        or root_dir.lower().endswith(".mov")
    ):
        video = np.asarray(mimread(root_dir))
        if len(video.shape) == 3:
            video = np.asarray([gray2rgb(frame) for frame in video])
        if video.shape[-1] == 4:
            video = video[..., :3]
        video_array = img_as_float32(video)
    else:
        raise Exception("Unknown file extensions  %s" % root_dir)

    return video_array

    class FramesDataset(Dataset):
        """
        Dataset of videos, each video can be represented as:
         - an image of concatenated frames
         - '.mp4' or '.gif'
         - folder with all frames
        """

    def __init__(
        self,
        root_dir,
        frame_shape=(384, 384, 3),
        id_sampling=False,
        is_train=True,
        random_seed=0,
        pairs_list=None,
        augmentation_params=None,
    ):
        self.root_dir = root_dir
        self.videos = os.listdir(root_dir)
        self.frame_shape = tuple(frame_shape)
        self.pairs_list = pairs_list
        self.id_sampling = id_sampling
        if os.path.exists(os.path.join(root_dir, "train")):
            assert os.path.exists(os.path.join(root_dir, "test"))
            print("Use predefined train-test split.")
            if id_sampling:
                train_videos = {
                    os.path.basename(video).split("#")[0]
                    for video in os.listdir(os.path.join(root_dir, "train"))
                }
                train_videos = list(train_videos)
            else:
                train_videos = os.listdir(os.path.join(root_dir, "train"))
            test_videos = os.listdir(os.path.join(root_dir, "test"))
            self.root_dir = os.path.join(self.root_dir, "train" if is_train else "test")
        else:
            print("Use random train-test split.")
            train_videos, test_videos = train_test_split(
                self.videos, random_state=random_seed, test_size=0.2
            )

        if is_train:
            self.videos = train_videos
        else:
            self.videos = test_videos

        self.is_train = is_train

        if self.is_train:
            self.transform = AllAugmentationTransform(**augmentation_params)
        else:
            self.transform = None

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

    def __getitem__(self, idx):
        if self.is_train and self.id_sampling:
            name = self.videos[idx]
            path = np.random.choice(
                glob.glob(os.path.join(self.root_dir, name + "*.mp4"))
            )
        else:
            name = self.videos[idx]
            path = os.path.join(self.root_dir, name)

        video_name = os.path.basename(path)

        if self.is_train and os.path.isdir(path):
            frames = os.listdir(path)
            num_frames = len(frames)
            frame_idx = np.sort(np.random.choice(num_frames, replace=True, size=2))
            video_array = [
                img_as_float32(io.imread(os.path.join(path, frames[idx])))
                for idx in frame_idx
            ]
        else:
            video_array = read_video(path, frame_shape=self.frame_shape)
            num_frames = len(video_array)
            frame_idx = (
                np.sort(np.random.choice(num_frames, replace=True, size=2))
                if self.is_train
                else range(num_frames)
            )
            video_array = video_array[frame_idx]

        if self.transform is not None:
            video_array = self.transform(video_array)

        out = {}
        if self.is_train:
            source = np.array(video_array[0], dtype="float32")
            driving = np.array(video_array[1], dtype="float32")

            out["driving"] = driving.transpose((2, 0, 1))
            out["source"] = source.transpose((2, 0, 1))
        else:
            video = np.array(video_array, dtype="float32")
            out["video"] = video.transpose((3, 0, 1, 2))

        out["name"] = video_name

        return out

    class DatasetRepeater(Dataset):
        """
        Pass several times over the same dataset for better i/o performance
        """

    def __init__(self, dataset, num_repeats=100):
        self.dataset = dataset
        self.num_repeats = num_repeats

    def __len__(self):
        return self.num_repeats * self.dataset.__len__()

    def __getitem__(self, idx):
        return self.dataset[idx % self.dataset.__len__()]

    class PairedDataset(Dataset):
        """
        Dataset of pairs for animation.
        """

    def __init__(self, initial_dataset, number_of_pairs, seed=0):
        self.initial_dataset = initial_dataset
        pairs_list = self.initial_dataset.pairs_list

        np.random.seed(seed)

        if pairs_list is None:
            max_idx = min(number_of_pairs, len(initial_dataset))
            nx, ny = max_idx, max_idx
            xy = np.mgrid[:nx, :ny].reshape(2, -1).T
            number_of_pairs = min(xy.shape[0], number_of_pairs)
            self.pairs = xy.take(
                np.random.choice(xy.shape[0], number_of_pairs, replace=False), axis=0
            )
        else:
            videos = self.initial_dataset.videos
            name_to_index = {name: index for index, name in enumerate(videos)}
            pairs = pd.read_csv(pairs_list)
            pairs = pairs[
                np.logical_and(
                    pairs["source"].isin(videos), pairs["driving"].isin(videos)
                )
            ]

            number_of_pairs = min(pairs.shape[0], number_of_pairs)
            self.pairs = []
            self.start_frames = []
            for ind in range(number_of_pairs):
                self.pairs.append(
                    (
                        name_to_index[pairs["driving"].iloc[ind]],
                        name_to_index[pairs["source"].iloc[ind]],
                    )
                )

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

    def __getitem__(self, idx):
        pair = self.pairs[idx]
        first = self.initial_dataset[pair[0]]
        second = self.initial_dataset[pair[1]]
        first = {"driving_" + key: value for key, value in first.items()}
        second = {"source_" + key: value for key, value in second.items()}

        return {**first, **second}