Rocky1/SadTalker

1import os2 3from tqdm import tqdm4import torch5import numpy as np6import random7import scipy.io as scio8import src.utils.audio as audio9 10def crop_pad_audio(wav, audio_length):11    if len(wav) > audio_length:12        wav = wav[:audio_length]13    elif len(wav) < audio_length:14        wav = np.pad(wav, [0, audio_length - len(wav)], mode='constant', constant_values=0)15    return wav16 17def parse_audio_length(audio_length, sr, fps):18    bit_per_frames = sr / fps19 20    num_frames = int(audio_length / bit_per_frames)21    audio_length = int(num_frames * bit_per_frames)22 23    return audio_length, num_frames24 25def generate_blink_seq(num_frames):26    ratio = np.zeros((num_frames,1))27    frame_id = 028    while frame_id in range(num_frames):29        start = 8030        if frame_id+start+9<=num_frames - 1:31            ratio[frame_id+start:frame_id+start+9, 0] = [0.5,0.6,0.7,0.9,1, 0.9, 0.7,0.6,0.5]32            frame_id = frame_id+start+933        else:34            break35    return ratio 36 37def generate_blink_seq_randomly(num_frames):38    ratio = np.zeros((num_frames,1))39    if num_frames<=20:40        return ratio41    frame_id = 042    while frame_id in range(num_frames):43        start = random.choice(range(min(10,num_frames), min(int(num_frames/2), 70))) 44        if frame_id+start+5<=num_frames - 1:45            ratio[frame_id+start:frame_id+start+5, 0] = [0.5, 0.9, 1.0, 0.9, 0.5]46            frame_id = frame_id+start+547        else:48            break49    return ratio50 51def get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path, still=False):52 53    syncnet_mel_step_size = 1654    fps = 2555 56    pic_name = os.path.splitext(os.path.split(first_coeff_path)[-1])[0]57    audio_name = os.path.splitext(os.path.split(audio_path)[-1])[0]58 59    wav = audio.load_wav(audio_path, 16000) 60    wav_length, num_frames = parse_audio_length(len(wav), 16000, 25)61    wav = crop_pad_audio(wav, wav_length)62    orig_mel = audio.melspectrogram(wav).T63    spec = orig_mel.copy()         # nframes 8064    indiv_mels = []65 66    for i in tqdm(range(num_frames), 'mel:'):67        start_frame_num = i-268        start_idx = int(80. * (start_frame_num / float(fps)))69        end_idx = start_idx + syncnet_mel_step_size70        seq = list(range(start_idx, end_idx))71        seq = [ min(max(item, 0), orig_mel.shape[0]-1) for item in seq ]72        m = spec[seq, :]73        indiv_mels.append(m.T)74    indiv_mels = np.asarray(indiv_mels)         # T 80 1675 76    ratio = generate_blink_seq_randomly(num_frames)      # T77    source_semantics_path = first_coeff_path78    source_semantics_dict = scio.loadmat(source_semantics_path)79    ref_coeff = source_semantics_dict['coeff_3dmm'][:1,:70]         #1 7080    ref_coeff = np.repeat(ref_coeff, num_frames, axis=0)81 82    if ref_eyeblink_coeff_path is not None:83        ratio[:num_frames] = 084        refeyeblink_coeff_dict = scio.loadmat(ref_eyeblink_coeff_path)85        refeyeblink_coeff = refeyeblink_coeff_dict['coeff_3dmm'][:,:64]86        refeyeblink_num_frames = refeyeblink_coeff.shape[0]87        if refeyeblink_num_frames<num_frames:88            div = num_frames//refeyeblink_num_frames89            re = num_frames%refeyeblink_num_frames90            refeyeblink_coeff_list = [refeyeblink_coeff for i in range(div)]91            refeyeblink_coeff_list.append(refeyeblink_coeff[:re, :64])92            refeyeblink_coeff = np.concatenate(refeyeblink_coeff_list, axis=0)93            print(refeyeblink_coeff.shape[0])94 95        ref_coeff[:, :64] = refeyeblink_coeff[:num_frames, :64] 96    97    indiv_mels = torch.FloatTensor(indiv_mels).unsqueeze(1).unsqueeze(0) # bs T 1 80 1698 99    if still:100        ratio = torch.FloatTensor(ratio).unsqueeze(0).fill_(0.)                        # bs T101    else:102        ratio = torch.FloatTensor(ratio).unsqueeze(0)103                               # bs T104    ref_coeff = torch.FloatTensor(ref_coeff).unsqueeze(0)                # bs 1 70105 106    indiv_mels = indiv_mels.to(device)107    ratio = ratio.to(device)108    ref_coeff = ref_coeff.to(device)109 110    return {'indiv_mels': indiv_mels,  111            'ref': ref_coeff, 112            'num_frames': num_frames, 113            'ratio_gt': ratio,114            'audio_name': audio_name, 'pic_name': pic_name}115 116