gulabpatel/First-Order-Motion

1import os2from skimage import io, img_as_float323from skimage.color import gray2rgb4from sklearn.model_selection import train_test_split5from imageio import mimread6 7import numpy as np8from torch.utils.data import Dataset9import pandas as pd10from augmentation import AllAugmentationTransform11import glob12 13 14def read_video(name, frame_shape):15    """16    Read video which can be:17      - an image of concatenated frames18      - '.mp4' and'.gif'19      - folder with videos20    """21 22    if os.path.isdir(name):23        frames = sorted(os.listdir(name))24        num_frames = len(frames)25        video_array = np.array(26            [img_as_float32(io.imread(os.path.join(name, frames[idx]))) for idx in range(num_frames)])27    elif name.lower().endswith('.png') or name.lower().endswith('.jpg'):28        image = io.imread(name)29 30        if len(image.shape) == 2 or image.shape[2] == 1:31            image = gray2rgb(image)32 33        if image.shape[2] == 4:34            image = image[..., :3]35 36        image = img_as_float32(image)37 38        video_array = np.moveaxis(image, 1, 0)39 40        video_array = video_array.reshape((-1,) + frame_shape)41        video_array = np.moveaxis(video_array, 1, 2)42    elif name.lower().endswith('.gif') or name.lower().endswith('.mp4') or name.lower().endswith('.mov'):43        video = np.array(mimread(name))44        if len(video.shape) == 3:45            video = np.array([gray2rgb(frame) for frame in video])46        if video.shape[-1] == 4:47            video = video[..., :3]48        video_array = img_as_float32(video)49    else:50        raise Exception("Unknown file extensions  %s" % name)51 52    return video_array53 54 55class FramesDataset(Dataset):56    """57    Dataset of videos, each video can be represented as:58      - an image of concatenated frames59      - '.mp4' or '.gif'60      - folder with all frames61    """62 63    def __init__(self, root_dir, frame_shape=(256, 256, 3), id_sampling=False, is_train=True,64                 random_seed=0, pairs_list=None, augmentation_params=None):65        self.root_dir = root_dir66        self.videos = os.listdir(root_dir)67        self.frame_shape = tuple(frame_shape)68        self.pairs_list = pairs_list69        self.id_sampling = id_sampling70        if os.path.exists(os.path.join(root_dir, 'train')):71            assert os.path.exists(os.path.join(root_dir, 'test'))72            print("Use predefined train-test split.")73            if id_sampling:74                train_videos = {os.path.basename(video).split('#')[0] for video in75                                os.listdir(os.path.join(root_dir, 'train'))}76                train_videos = list(train_videos)77            else:78                train_videos = os.listdir(os.path.join(root_dir, 'train'))79            test_videos = os.listdir(os.path.join(root_dir, 'test'))80            self.root_dir = os.path.join(self.root_dir, 'train' if is_train else 'test')81        else:82            print("Use random train-test split.")83            train_videos, test_videos = train_test_split(self.videos, random_state=random_seed, test_size=0.2)84 85        if is_train:86            self.videos = train_videos87        else:88            self.videos = test_videos89 90        self.is_train = is_train91 92        if self.is_train:93            self.transform = AllAugmentationTransform(**augmentation_params)94        else:95            self.transform = None96 97    def __len__(self):98        return len(self.videos)99 100    def __getitem__(self, idx):101        if self.is_train and self.id_sampling:102            name = self.videos[idx]103            path = np.random.choice(glob.glob(os.path.join(self.root_dir, name + '*.mp4')))104        else:105            name = self.videos[idx]106            path = os.path.join(self.root_dir, name)107 108        video_name = os.path.basename(path)109 110        if self.is_train and os.path.isdir(path):111            frames = os.listdir(path)112            num_frames = len(frames)113            frame_idx = np.sort(np.random.choice(num_frames, replace=True, size=2))114            video_array = [img_as_float32(io.imread(os.path.join(path, frames[idx]))) for idx in frame_idx]115        else:116            video_array = read_video(path, frame_shape=self.frame_shape)117            num_frames = len(video_array)118            frame_idx = np.sort(np.random.choice(num_frames, replace=True, size=2)) if self.is_train else range(119                num_frames)120            video_array = video_array[frame_idx]121 122        if self.transform is not None:123            video_array = self.transform(video_array)124 125        out = {}126        if self.is_train:127            source = np.array(video_array[0], dtype='float32')128            driving = np.array(video_array[1], dtype='float32')129 130            out['driving'] = driving.transpose((2, 0, 1))131            out['source'] = source.transpose((2, 0, 1))132        else:133            video = np.array(video_array, dtype='float32')134            out['video'] = video.transpose((3, 0, 1, 2))135 136        out['name'] = video_name137 138        return out139 140 141class DatasetRepeater(Dataset):142    """143    Pass several times over the same dataset for better i/o performance144    """145 146    def __init__(self, dataset, num_repeats=100):147        self.dataset = dataset148        self.num_repeats = num_repeats149 150    def __len__(self):151        return self.num_repeats * self.dataset.__len__()152 153    def __getitem__(self, idx):154        return self.dataset[idx % self.dataset.__len__()]155 156 157class PairedDataset(Dataset):158    """159    Dataset of pairs for animation.160    """161 162    def __init__(self, initial_dataset, number_of_pairs, seed=0):163        self.initial_dataset = initial_dataset164        pairs_list = self.initial_dataset.pairs_list165 166        np.random.seed(seed)167 168        if pairs_list is None:169            max_idx = min(number_of_pairs, len(initial_dataset))170            nx, ny = max_idx, max_idx171            xy = np.mgrid[:nx, :ny].reshape(2, -1).T172            number_of_pairs = min(xy.shape[0], number_of_pairs)173            self.pairs = xy.take(np.random.choice(xy.shape[0], number_of_pairs, replace=False), axis=0)174        else:175            videos = self.initial_dataset.videos176            name_to_index = {name: index for index, name in enumerate(videos)}177            pairs = pd.read_csv(pairs_list)178            pairs = pairs[np.logical_and(pairs['source'].isin(videos), pairs['driving'].isin(videos))]179 180            number_of_pairs = min(pairs.shape[0], number_of_pairs)181            self.pairs = []182            self.start_frames = []183            for ind in range(number_of_pairs):184                self.pairs.append(185                    (name_to_index[pairs['driving'].iloc[ind]], name_to_index[pairs['source'].iloc[ind]]))186 187    def __len__(self):188        return len(self.pairs)189 190    def __getitem__(self, idx):191        pair = self.pairs[idx]192        first = self.initial_dataset[pair[0]]193        second = self.initial_dataset[pair[1]]194        first = {'driving_' + key: value for key, value in first.items()}195        second = {'source_' + key: value for key, value in second.items()}196 197        return {**first, **second}198