thaint2901/talking-head-generation-deploy

1import os2from skimage import io, img_as_float323from skimage.color import gray2rgb4from sklearn.model_selection import train_test_split5from imageio import mimread6from skimage.transform import resize7import numpy as np8from torch.utils.data import Dataset9from augmentation import AllAugmentationTransform10import glob11from functools import partial12 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 = mimread(name)44        if len(video[0].shape) == 2:45            video = [gray2rgb(frame) for frame in video]46        if frame_shape is not None:47            video = np.array([resize(frame, frame_shape) for frame in video])48        video = np.array(video)49        if video.shape[-1] == 4:50            video = video[..., :3]51        video_array = img_as_float32(video)52    else:53        raise Exception("Unknown file extensions  %s" % name)54 55    return video_array56 57 58class FramesDataset(Dataset):59    """60    Dataset of videos, each video can be represented as:61      - an image of concatenated frames62      - '.mp4' or '.gif'63      - folder with all frames64    """65 66    def __init__(self, root_dir, frame_shape=(256, 256, 3), id_sampling=False, is_train=True,67                 random_seed=0, pairs_list=None, augmentation_params=None):68        self.root_dir = root_dir69        self.videos = os.listdir(root_dir)70        self.frame_shape = frame_shape71        print(self.frame_shape)72        self.pairs_list = pairs_list73        self.id_sampling = id_sampling74 75        if os.path.exists(os.path.join(root_dir, 'train')):76            assert os.path.exists(os.path.join(root_dir, 'test'))77            print("Use predefined train-test split.")78            if id_sampling:79                train_videos = {os.path.basename(video).split('#')[0] for video in80                                os.listdir(os.path.join(root_dir, 'train'))}81                train_videos = list(train_videos)82            else:83                train_videos = os.listdir(os.path.join(root_dir, 'train'))84            test_videos = os.listdir(os.path.join(root_dir, 'test'))85            self.root_dir = os.path.join(self.root_dir, 'train' if is_train else 'test')86        else:87            print("Use random train-test split.")88            train_videos, test_videos = train_test_split(self.videos, random_state=random_seed, test_size=0.2)89 90        if is_train:91            self.videos = train_videos92        else:93            self.videos = test_videos94 95        self.is_train = is_train96 97        if self.is_train:98            self.transform = AllAugmentationTransform(**augmentation_params)99        else:100            self.transform = None101 102    def __len__(self):103        return len(self.videos)104 105    def __getitem__(self, idx):106            107        if self.is_train and self.id_sampling:   108            name = self.videos[idx]109            path = np.random.choice(glob.glob(os.path.join(self.root_dir, name + '*.mp4')))110        else:111            name = self.videos[idx]112            path = os.path.join(self.root_dir, name)113 114        video_name = os.path.basename(path)115        if self.is_train and os.path.isdir(path):116 117            frames = os.listdir(path)118            num_frames = len(frames)119            frame_idx = np.sort(np.random.choice(num_frames, replace=True, size=2))120            121            if self.frame_shape is not None:122                resize_fn = partial(resize, output_shape=self.frame_shape)123            else:124                resize_fn = img_as_float32125 126            if type(frames[0]) is bytes:127                video_array = [resize_fn(io.imread(os.path.join(path, frames[idx].decode('utf-8')))) for idx in128                            frame_idx]129            else:130                video_array = [resize_fn(io.imread(os.path.join(path, frames[idx]))) for idx in frame_idx]131        else:132                 133            video_array = read_video(path, frame_shape=self.frame_shape)134            135            num_frames = len(video_array)136            frame_idx = np.sort(np.random.choice(num_frames, replace=True, size=2)) if self.is_train else range(137                num_frames)138            video_array = video_array[frame_idx]139            140 141        if self.transform is not None:142            video_array = self.transform(video_array)143 144        out = {}145        if self.is_train:146            source = np.array(video_array[0], dtype='float32')147            driving = np.array(video_array[1], dtype='float32')148 149            out['driving'] = driving.transpose((2, 0, 1))150            out['source'] = source.transpose((2, 0, 1))151        else:152            video = np.array(video_array, dtype='float32')153            out['video'] = video.transpose((3, 0, 1, 2))154 155        out['name'] = video_name156        return out157 158 159class DatasetRepeater(Dataset):160    """161    Pass several times over the same dataset for better i/o performance162    """163 164    def __init__(self, dataset, num_repeats=100):165        self.dataset = dataset166        self.num_repeats = num_repeats167 168    def __len__(self):169        return self.num_repeats * self.dataset.__len__()170 171    def __getitem__(self, idx):172        return self.dataset[idx % self.dataset.__len__()]173 174