Spaces:
Runtime error
Runtime error
| import os | |
| from skimage import io, img_as_float32, transform | |
| 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 | |
| import pandas as pd | |
| from augmentation import AllAugmentationTransform | |
| import glob | |
| import pickle | |
| import random | |
| def read_video(name, frame_shape): | |
| """ | |
| Read video which can be: | |
| - an image of concatenated frames | |
| - '.mp4' and'.gif' | |
| - folder with videos | |
| """ | |
| if os.path.isdir(name): | |
| frames = sorted(os.listdir(name)) | |
| num_frames = len(frames) | |
| video_array = np.array( | |
| [img_as_float32(io.imread(os.path.join(name, frames[idx]))) for idx in range(num_frames)]) | |
| elif name.lower().endswith('.png') or name.lower().endswith('.jpg'): | |
| image = io.imread(name) | |
| 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 = video_array.reshape((-1,) + frame_shape) | |
| video_array = np.moveaxis(video_array, 1, 2) | |
| elif name.lower().endswith('.gif') or name.lower().endswith('.mp4') or name.lower().endswith('.mov'): | |
| video = np.array(mimread(name)) | |
| if len(video.shape) == 3: | |
| video = np.array([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" % name) | |
| return video_array | |
| def get_list(ipath,base_name): | |
| #ipath = '/mnt/lustre/share/jixinya/LRW/pose/train_fo/' | |
| ipath = os.path.join(ipath,base_name) | |
| name_list = os.listdir(ipath) | |
| image_path = os.path.join('/mnt/lustre/share/jixinya/LRW/Image/',base_name) | |
| all = [] | |
| for k in range(len(name_list)): | |
| name = name_list[k] | |
| path_ = os.path.join(ipath,name) | |
| Dir = os.listdir(path_) | |
| for i in range(len(Dir)): | |
| word = Dir[i] | |
| path = os.path.join(path_, word) | |
| if os.path.exists(os.path.join(image_path,name,word.split('.')[0])): | |
| all.append(name+'/'+word.split('.')[0]) | |
| #print(k,name,i,word) | |
| print('get list '+os.path.basename(ipath)) | |
| return all | |
| class AudioDataset(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=(256, 256, 3), id_sampling=False, is_train=True, | |
| random_seed=0, pairs_list=None, augmentation_params=None): | |
| self.root_dir = root_dir | |
| self.audio_dir = os.path.join(root_dir,'MFCC') | |
| self.image_dir = os.path.join(root_dir,'Image') | |
| self.landmark_dir = os.path.join(root_dir,'Landmark') | |
| self.pose_dir = os.path.join(root_dir,'pose') | |
| # assert len(os.listdir(self.audio_dir)) == len(os.listdir(self.image_dir)), 'audio and image length not equal' | |
| df=open('../LRW/list/test_fo.txt','rb') | |
| self.videos=pickle.load(df) | |
| df.close() | |
| # self.videos=np.load('../LRW/list/train_fo.npy') | |
| # self.videos = os.listdir(self.landmark_dir) | |
| self.frame_shape = tuple(frame_shape) | |
| self.pairs_list = pairs_list | |
| self.id_sampling = id_sampling | |
| self.pca = np.load('../LRW/list/U_106.npy')[:, :16] | |
| self.mean = np.load('../LRW/list/mean_106.npy') | |
| if os.path.exists(os.path.join(self.pose_dir, 'train_fo')): | |
| assert os.path.exists(os.path.join(self.pose_dir, 'test_fo')) | |
| 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(self.image_dir, 'train'))} | |
| train_videos = list(train_videos) | |
| else: | |
| train_videos = np.load('../LRW/list/train_fo.npy')# get_list(self.pose_dir, 'train_fo') | |
| df=open('../LRW/list/test_fo.txt','rb') | |
| test_videos=pickle.load(df) | |
| df.close() | |
| # test_videos = np.load('../LRW/list/train_fo.npy') | |
| #get_list(self.pose_dir, 'test_fo') | |
| # self.root_dir = os.path.join(self.root_dir, 'train' if is_train else 'test') | |
| self.landmark_dir = os.path.join(self.landmark_dir, 'train_fo' if is_train else 'test_fo') | |
| self.image_dir = os.path.join(self.image_dir, 'train_fo' if is_train else 'test_fo') | |
| self.audio_dir = os.path.join(self.audio_dir, 'train' if is_train else 'test') | |
| self.pose_dir = os.path.join(self.pose_dir, 'train_fo' if is_train else 'test_fo') | |
| 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].split('.')[0] | |
| path = np.random.choice(glob.glob(os.path.join(self.root_dir, name + '*.mp4'))) | |
| else: | |
| name = self.videos[idx].split('.')[0] | |
| landmark_path = os.path.join(self.landmark_dir, name+'.npy') | |
| audio_path = os.path.join(self.audio_dir, name) | |
| pose_path = os.path.join(self.pose_dir,name) | |
| path = os.path.join(self.image_dir, name) | |
| video_name = os.path.basename(path) | |
| if os.path.isdir(path): | |
| # if self.is_train and os.path.isdir(path): | |
| lmark = np.load(landmark_path).reshape(-1,212)/255 | |
| if np.isnan(lmark).sum() or np.isinf(lmark).sum(): | |
| print('Wrong lmark '+ video_name, file=open('log/wrong.txt', 'a')) | |
| lmark = np.zeros((29,212)) | |
| lmark = lmark - self.mean | |
| lmark = np.dot(lmark, self.pca) | |
| # mfcc loading | |
| r = random.choice([x for x in range(3, 8)]) | |
| example_landmark = lmark[r, :] | |
| example_image = img_as_float32(io.imread(os.path.join(path, str(r)+'.png'))) | |
| # example_mfcc = mfcc[(r - 3) * 4: (r + 4) * 4, 1:] | |
| mfccs = [] | |
| for ind in range(1, 17): | |
| # t_mfcc = mfcc[(r + ind - 3) * 4: (r + ind + 4) * 4, 1:] | |
| try: | |
| t_mfcc = np.load(os.path.join(audio_path,str(r + ind)+'.npy'),allow_pickle=True)[:, 1:] | |
| if np.isnan(t_mfcc).sum() or np.isinf(t_mfcc).sum(): | |
| print('Wrong mfcc '+ video_name+str(r+ind), file=open('log/wrong.txt', 'a')) | |
| t_mfcc = np.zeros((28,13))[:,1:] | |
| except: | |
| t_mfcc = np.zeros((28,13))[:,1:] | |
| mfccs.append(t_mfcc) | |
| mfccs = np.array(mfccs) | |
| if not self.is_train: | |
| poses = [] | |
| video_array = [] | |
| for ind in range(1, 17): | |
| # t_mfcc = mfcc[(r + ind - 3) * 4: (r + ind + 4) * 4, 1:] | |
| t_pose = np.load(os.path.join(pose_path,str(r + ind)+'.npy'))[:-1] | |
| poses.append(t_pose) | |
| image = img_as_float32(io.imread(os.path.join(path, str(r + ind)+'.png'))) | |
| video_array.append(image) | |
| poses = np.array(poses) | |
| video_array = np.array(video_array) | |
| else: | |
| poses = [] | |
| video_array = [] | |
| for ind in range(1, 17): | |
| # t_mfcc = mfcc[(r + ind - 3) * 4: (r + ind + 4) * 4, 1:] | |
| t_pose = np.load(os.path.join(self.pose_dir,name+'.npy'))[r+ind,:-1] | |
| if np.isnan(t_pose).sum() or np.isinf(t_pose).sum(): | |
| print('Wrong pose '+ video_name, file=open('log/wrong.txt', 'a')) | |
| t_pose = np.zeros((6,)) | |
| poses.append(t_pose) | |
| image = img_as_float32(io.imread(os.path.join(path, str(r + ind)+'.png'))) | |
| video_array.append(image) | |
| poses = np.array(poses) | |
| video_array = np.array(video_array) | |
| #mfccs = torch.FloatTensor(mfccs) | |
| landmark = lmark[r + 1: r + 17, :] | |
| index_32 = [0,4,8,12,16,20,24,28,32,33,35,67,68,40,42,52,55,72,73,58,61,75,76,46,47,51,84,87,90,93,98,102] | |
| driving_landmark = np.load(landmark_path)[r + 1: r + 17, :][:,index_32] | |
| source_landmark = np.load(landmark_path)[r, :][index_32] | |
| 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 True:#self.is_train: | |
| # a = img_as_float32(io.imread('/media/thea/Data/first-order-model/images_512/102.jpg')) | |
| # source = np.array(a, dtype='float32') | |
| driving = np.array(video_array, dtype='float32') | |
| spatial_size = np.array(driving.shape[1:3][::-1])[np.newaxis] | |
| # example_landmark = np.array(2*example_landmark / spatial_size -1, dtype='float32') | |
| driving_landmark = np.array(2*driving_landmark / spatial_size -1, dtype='float32') | |
| source_landmark = np.array(2*source_landmark / spatial_size -1, dtype='float32') | |
| driving_pose = np.array(poses, dtype='float32') | |
| example_landmark = np.array(example_landmark, dtype='float32') | |
| example_image = np.array(example_image, dtype='float32') | |
| # source_cube = np.array(transform.resize(cube_array[0], (64,64)), dtype='float32') | |
| # driving_cube = np.array(transform.resize(cube_array[1], (64,64)), dtype='float32') | |
| # source_heatmap = np.array(heatmap_array[0] , dtype='float32') | |
| # driving_heatmap = np.array(heatmap_array[1] , dtype='float32') | |
| # out['source_cube'] = source_cube | |
| # out['driving_cube'] = driving_cube | |
| out['example_landmark'] = example_landmark | |
| out['example_image'] = example_image.transpose((2, 0, 1)) | |
| out['driving_landmark'] = driving_landmark | |
| out['source_landmark'] = source_landmark | |
| out['driving_pose'] = driving_pose | |
| # out['source_heatmap'] = source_heatmap | |
| # out['driving_heatmap'] = driving_heatmap | |
| out['driving'] = driving.transpose((0, 3, 1, 2)) | |
| # out['source'] = source.transpose((2, 0, 1)) | |
| # out['source_audio'] = np.array(audio_array[0], dtype='float32') | |
| out['driving_audio'] = np.array(mfccs, dtype='float32') | |
| out['gt_landmark'] = np.array(landmark, dtype='float32') | |
| out['pca'] = np.array(self.pca, dtype='float32') | |
| out['mean'] = np.array(self.mean, dtype='float32') | |
| out['name'] = video_name | |
| return out | |
| 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=(256, 256, 3), id_sampling=False, is_train=True, | |
| random_seed=0, pairs_list=None, augmentation_params=None): | |
| self.root_dir = root_dir | |
| self.audio_dir = os.path.join(root_dir,'audio/') | |
| self.image_dir = os.path.join(root_dir,'image/') | |
| self.landmark_dir = os.path.join(root_dir,'cube/') | |
| # assert len(os.listdir(self.audio_dir)) == len(os.listdir(self.image_dir)), 'audio and image length not equal' | |
| df=open('/media/thea/新加卷/MEAD/neutral/train.txt','rb') | |
| self.videos=pickle.load(df) | |
| df.close() | |
| # self.videos = os.listdir(self.landmark_dir) | |
| self.frame_shape = tuple(frame_shape) | |
| self.pairs_list = pairs_list | |
| self.id_sampling = id_sampling | |
| if os.path.exists(os.path.join(self.image_dir, 'train')): | |
| assert os.path.exists(os.path.join(self.image_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(self.image_dir, 'train'))} | |
| train_videos = list(train_videos) | |
| else: | |
| train_videos = os.listdir(os.path.join(self.image_dir, 'train')) | |
| test_videos = os.listdir(os.path.join(self.image_dir, 'test')) | |
| self.root_dir = os.path.join(self.root_dir, 'train' if is_train else 'test') | |
| self.landmark_dir = os.path.join(self.landmark_dir, 'train' if is_train else 'test') | |
| self.image_dir = os.path.join(self.image_dir, 'train' if is_train else 'test') | |
| self.audio_dir = os.path.join(self.audio_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].split('.')[0] | |
| path = np.random.choice(glob.glob(os.path.join(self.root_dir, name + '*.mp4'))) | |
| else: | |
| name = self.videos[idx].split('.')[0] | |
| landmark_path = os.path.join(self.landmark_dir, name) | |
| audio_path = os.path.join(self.audio_dir, name) | |
| path = os.path.join(self.image_dir, name) | |
| video_name = os.path.basename(path) | |
| if self.is_train and os.path.isdir(path): | |
| frames = os.listdir(audio_path) | |
| num_frames = len(frames) | |
| frame_idx = np.sort(np.random.choice(num_frames-1, replace=True, size=2)) | |
| # landmark = np.load(landmark_path)#+'.npy' | |
| # assert len(os.listdir(path)) == len(landmark), video_name+' length not equal' | |
| video_array = [img_as_float32(io.imread(os.path.join(path, str(idx)+'.png'))) for idx in frame_idx] | |
| cube_array = [img_as_float32(io.imread(os.path.join(landmark_path, str(idx)+'.jpg'))) for idx in frame_idx] | |
| audio_array = [np.load(os.path.join(audio_path, str(idx)+'.npy'))[:,1:] for idx in frame_idx] | |
| index_20 = [0,16,32,35,40,52,55,58,61,46,72,73,75,76,84,87,90,93,98,102] | |
| index_32 = [0,4,8,12,16,20,24,28,32,33,35,67,68,40,42,52,55,72,73,58,61,75,76,46,47,51,84,87,90,93,98,102] | |
| # landmark_array = [landmark[idx] for idx in frame_idx] | |
| # landmark_array = [landmark[idx][index_32] 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: | |
| # a = img_as_float32(io.imread('/media/thea/Data/first-order-model/images_512/102.jpg')) | |
| # source = np.array(a, dtype='float32') | |
| source = np.array(video_array[0], dtype='float32') | |
| driving = np.array(video_array[1], dtype='float32') | |
| spatial_size = np.array(source.shape[:2][::-1])[np.newaxis] | |
| # source_landmark = np.array(2*landmark_array[0] / spatial_size -1, dtype='float32') | |
| # driving_landmark = np.array(2*landmark_array[1] / spatial_size -1, dtype='float32') | |
| source_cube = np.array(transform.resize(cube_array[0], (64,64)), dtype='float32') | |
| driving_cube = np.array(transform.resize(cube_array[1], (64,64)), dtype='float32') | |
| # source_heatmap = np.array(heatmap_array[0] , dtype='float32') | |
| # driving_heatmap = np.array(heatmap_array[1] , dtype='float32') | |
| out['source_cube'] = source_cube | |
| out['driving_cube'] = driving_cube | |
| # out['source_landmark'] = source_landmark | |
| # out['driving_landmark'] = driving_landmark | |
| # out['source_heatmap'] = source_heatmap | |
| # out['driving_heatmap'] = driving_heatmap | |
| out['driving'] = driving.transpose((2, 0, 1)) | |
| out['source'] = source.transpose((2, 0, 1)) | |
| out['source_audio'] = np.array(audio_array[0], dtype='float32') | |
| out['driving_audio'] = np.array(audio_array[1], dtype='float32') | |
| 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__()]#% 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} | |