import os import librosa from PIL import Image from torchvision import transforms import python_speech_features import random import os import numpy as np from tqdm import tqdm import torchvision import torchvision.transforms as transforms from PIL import Image class LatentDataLoader(object): def __init__( self, window_size, frame_jpgs, lmd_feats_prefix, audio_prefix, raw_audio_prefix, motion_latents_prefix, pose_prefix, db_name, video_fps=25, audio_hz=50, size=256, mfcc_mode=False, ): self.window_size = window_size self.lmd_feats_prefix = lmd_feats_prefix self.audio_prefix = audio_prefix self.pose_prefix = pose_prefix self.video_fps = video_fps self.audio_hz = audio_hz self.db_name = db_name self.raw_audio_prefix = raw_audio_prefix self.mfcc_mode = mfcc_mode self.transform = torchvision.transforms.Compose([ transforms.Resize((size, size)), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))] ) self.data = [] for db_name in [ 'VoxCeleb2', 'HDTF' ]: db_png_path = os.path.join(frame_jpgs, db_name) for clip_name in tqdm(os.listdir(db_png_path)): item_dict = dict() item_dict['clip_name'] = clip_name item_dict['frame_count'] = len(list(os.listdir(os.path.join(frame_jpgs, db_name, clip_name)))) item_dict['hubert_path'] = os.path.join(audio_prefix, db_name, clip_name +".npy") item_dict['wav_path'] = os.path.join(raw_audio_prefix, db_name, clip_name +".wav") item_dict['yaw_pitch_roll_path'] = os.path.join(pose_prefix, db_name, 'raw_videos_pose_yaw_pitch_roll', clip_name +".npy") if not os.path.exists(item_dict['yaw_pitch_roll_path']): print(f"{db_name}'s {clip_name} miss yaw_pitch_roll_path") continue item_dict['yaw_pitch_roll'] = np.load(item_dict['yaw_pitch_roll_path']) item_dict['yaw_pitch_roll'] = np.clip(item_dict['yaw_pitch_roll'], -90, 90) / 90.0 if not os.path.exists(item_dict['wav_path']): print(f"{db_name}'s {clip_name} miss wav_path") continue if not os.path.exists(item_dict['hubert_path']): print(f"{db_name}'s {clip_name} miss hubert_path") continue if self.mfcc_mode: wav, sr = librosa.load(item_dict['wav_path'], sr=16000) input_values = python_speech_features.mfcc(signal=wav,samplerate=sr,numcep=13,winlen=0.025,winstep=0.01) d_mfcc_feat = python_speech_features.base.delta(input_values, 1) d_mfcc_feat2 = python_speech_features.base.delta(input_values, 2) input_values = np.hstack((input_values, d_mfcc_feat, d_mfcc_feat2)) item_dict['hubert_obj'] = input_values else: item_dict['hubert_obj'] = np.load(item_dict['hubert_path'], mmap_mode='r') item_dict['lmd_path'] = os.path.join(lmd_feats_prefix, db_name, clip_name +".txt") item_dict['lmd_obj_full'] = self.read_landmark_info(item_dict['lmd_path'], upper_face=False) motion_start_path = os.path.join(motion_latents_prefix, db_name, 'motions', clip_name +".npy") motion_direction_path = os.path.join(motion_latents_prefix, db_name, 'directions', clip_name +".npy") if not os.path.exists(motion_start_path): print(f"{db_name}'s {clip_name} miss motion_start_path") continue if not os.path.exists(motion_direction_path): print(f"{db_name}'s {clip_name} miss motion_direction_path") continue item_dict['motion_start_obj'] = np.load(motion_start_path) item_dict['motion_direction_obj'] = np.load(motion_direction_path) if self.mfcc_mode: min_len = min( item_dict['lmd_obj_full'].shape[0], item_dict['yaw_pitch_roll'].shape[0], item_dict['motion_start_obj'].shape[0], item_dict['motion_direction_obj'].shape[0], int(item_dict['hubert_obj'].shape[0]/4), item_dict['frame_count'] ) item_dict['frame_count'] = min_len item_dict['hubert_obj'] = item_dict['hubert_obj'][:min_len*4,:] else: min_len = min( item_dict['lmd_obj_full'].shape[0], item_dict['yaw_pitch_roll'].shape[0], item_dict['motion_start_obj'].shape[0], item_dict['motion_direction_obj'].shape[0], int(item_dict['hubert_obj'].shape[1]/2), item_dict['frame_count'] ) item_dict['frame_count'] = min_len item_dict['hubert_obj'] = item_dict['hubert_obj'][:, :min_len*2, :] if min_len < self.window_size * self.video_fps + 5: continue print('Db count:', len(self.data)) def get_single_image(self, image_path): img_source = Image.open(image_path).convert('RGB') img_source = self.transform(img_source) return img_source def get_multiple_ranges(self, lists, multi_ranges): # Ensure that multi_ranges is a list of tuples if not all(isinstance(item, tuple) and len(item) == 2 for item in multi_ranges): raise ValueError("multi_ranges must be a list of (start, end) tuples with exactly two elements each") extracted_elements = [lists[start:end] for start, end in multi_ranges] flat_list = [item for sublist in extracted_elements for item in sublist] return flat_list def read_landmark_info(self, lmd_path, upper_face=True): with open(lmd_path, 'r') as file: lmd_lines = file.readlines() lmd_lines.sort() total_lmd_obj = [] for i, line in enumerate(lmd_lines): # Split the coordinates and filter out any empty strings coords = [c for c in line.strip().split(' ') if c] coords = coords[1:] # do not include the file name in the first row lmd_obj = [] if upper_face: # Ensure that the coordinates are parsed as integers for coord_pair in self.get_multiple_ranges(coords, [(0, 3), (14, 27), (36, 48)]): # 28δΈͺ x, y = coord_pair.split('_') lmd_obj.append((int(x)/512, int(y)/512)) else: for coord_pair in coords: x, y = coord_pair.split('_') lmd_obj.append((int(x)/512, int(y)/512)) total_lmd_obj.append(lmd_obj) return np.array(total_lmd_obj, dtype=np.float32) def calculate_face_height(self, landmarks): forehead_center = (landmarks[ :, 21, :] + landmarks[:, 22, :]) / 2 chin_bottom = landmarks[:, 8, :] distances = np.linalg.norm(forehead_center - chin_bottom, axis=1, keepdims=True) return distances def __getitem__(self, index): data_item = self.data[index] hubert_obj = data_item['hubert_obj'] frame_count = data_item['frame_count'] lmd_obj_full = data_item['lmd_obj_full'] yaw_pitch_roll = data_item['yaw_pitch_roll'] motion_start_obj = data_item['motion_start_obj'] motion_direction_obj = data_item['motion_direction_obj'] frame_end_index = random.randint(self.window_size * self.video_fps + 1, frame_count - 1) frame_start_index = frame_end_index - self.window_size * self.video_fps frame_hint_index = frame_start_index - 1 audio_start_index = int(frame_start_index * (self.audio_hz / self.video_fps)) audio_end_index = int(frame_end_index * (self.audio_hz / self.video_fps)) if self.mfcc_mode: audio_feats = hubert_obj[audio_start_index:audio_end_index, :] else: audio_feats = hubert_obj[:, audio_start_index:audio_end_index, :] lmd_obj_full = lmd_obj_full[frame_hint_index:frame_end_index, :] yaw_pitch_roll = yaw_pitch_roll[frame_start_index:frame_end_index, :] motion_start = motion_start_obj[frame_hint_index] motion_direction_start = motion_direction_obj[frame_hint_index] motion_direction = motion_direction_obj[frame_start_index:frame_end_index, :] return { 'motion_start': motion_start, 'motion_direction': motion_direction, 'audio_feats': audio_feats, 'face_location': lmd_obj_full[1:, 30, 0], # '1:' means taking the first frame as the driven frame. '30' is the noise location, '0' means x coordinate 'face_scale': self.calculate_face_height(lmd_obj_full[1:,:,:]), 'yaw_pitch_roll': yaw_pitch_roll, 'motion_direction_start': motion_direction_start, } def __len__(self): return len(self.data)