mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-15 00:03:06 +00:00
Add files via upload
This commit is contained in:
Justin John
@@ -0,0 +1,120 @@
|
||||
import os
|
||||
|
||||
from tqdm import tqdm
|
||||
import torch
|
||||
import numpy as np
|
||||
import random
|
||||
import scipy.io as scio
|
||||
import src.utils.audio as audio
|
||||
|
||||
def crop_pad_audio(wav, audio_length):
|
||||
if len(wav) > audio_length:
|
||||
wav = wav[:audio_length]
|
||||
elif len(wav) < audio_length:
|
||||
wav = np.pad(wav, [0, audio_length - len(wav)], mode='constant', constant_values=0)
|
||||
return wav
|
||||
|
||||
def parse_audio_length(audio_length, sr, fps):
|
||||
bit_per_frames = sr / fps
|
||||
|
||||
num_frames = int(audio_length / bit_per_frames)
|
||||
audio_length = int(num_frames * bit_per_frames)
|
||||
|
||||
return audio_length, num_frames
|
||||
|
||||
def generate_blink_seq(num_frames):
|
||||
ratio = np.zeros((num_frames,1))
|
||||
frame_id = 0
|
||||
while frame_id in range(num_frames):
|
||||
start = 80
|
||||
if frame_id+start+9<=num_frames - 1:
|
||||
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]
|
||||
frame_id = frame_id+start+9
|
||||
else:
|
||||
break
|
||||
return ratio
|
||||
|
||||
def generate_blink_seq_randomly(num_frames):
|
||||
ratio = np.zeros((num_frames,1))
|
||||
if num_frames<=20:
|
||||
return ratio
|
||||
frame_id = 0
|
||||
while frame_id in range(num_frames):
|
||||
start = random.choice(range(min(10,num_frames), min(int(num_frames/2), 70)))
|
||||
if frame_id+start+5<=num_frames - 1:
|
||||
ratio[frame_id+start:frame_id+start+5, 0] = [0.5, 0.9, 1.0, 0.9, 0.5]
|
||||
frame_id = frame_id+start+5
|
||||
else:
|
||||
break
|
||||
return ratio
|
||||
|
||||
def get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path, still=False, idlemode=False, length_of_audio=False, use_blink=True):
|
||||
|
||||
syncnet_mel_step_size = 16
|
||||
fps = 25
|
||||
|
||||
pic_name = os.path.splitext(os.path.split(first_coeff_path)[-1])[0]
|
||||
audio_name = os.path.splitext(os.path.split(audio_path)[-1])[0]
|
||||
|
||||
|
||||
if idlemode:
|
||||
num_frames = int(length_of_audio * 25)
|
||||
indiv_mels = np.zeros((num_frames, 80, 16))
|
||||
else:
|
||||
wav = audio.load_wav(audio_path, 16000)
|
||||
wav_length, num_frames = parse_audio_length(len(wav), 16000, 25)
|
||||
wav = crop_pad_audio(wav, wav_length)
|
||||
orig_mel = audio.melspectrogram(wav).T
|
||||
spec = orig_mel.copy() # nframes 80
|
||||
indiv_mels = []
|
||||
|
||||
for i in tqdm(range(num_frames), 'mel:'):
|
||||
start_frame_num = i-2
|
||||
start_idx = int(80. * (start_frame_num / float(fps)))
|
||||
end_idx = start_idx + syncnet_mel_step_size
|
||||
seq = list(range(start_idx, end_idx))
|
||||
seq = [ min(max(item, 0), orig_mel.shape[0]-1) for item in seq ]
|
||||
m = spec[seq, :]
|
||||
indiv_mels.append(m.T)
|
||||
indiv_mels = np.asarray(indiv_mels) # T 80 16
|
||||
|
||||
ratio = generate_blink_seq_randomly(num_frames) # T
|
||||
source_semantics_path = first_coeff_path
|
||||
source_semantics_dict = scio.loadmat(source_semantics_path)
|
||||
ref_coeff = source_semantics_dict['coeff_3dmm'][:1,:70] #1 70
|
||||
ref_coeff = np.repeat(ref_coeff, num_frames, axis=0)
|
||||
|
||||
if ref_eyeblink_coeff_path is not None:
|
||||
ratio[:num_frames] = 0
|
||||
refeyeblink_coeff_dict = scio.loadmat(ref_eyeblink_coeff_path)
|
||||
refeyeblink_coeff = refeyeblink_coeff_dict['coeff_3dmm'][:,:64]
|
||||
refeyeblink_num_frames = refeyeblink_coeff.shape[0]
|
||||
if refeyeblink_num_frames<num_frames:
|
||||
div = num_frames//refeyeblink_num_frames
|
||||
re = num_frames%refeyeblink_num_frames
|
||||
refeyeblink_coeff_list = [refeyeblink_coeff for i in range(div)]
|
||||
refeyeblink_coeff_list.append(refeyeblink_coeff[:re, :64])
|
||||
refeyeblink_coeff = np.concatenate(refeyeblink_coeff_list, axis=0)
|
||||
print(refeyeblink_coeff.shape[0])
|
||||
|
||||
ref_coeff[:, :64] = refeyeblink_coeff[:num_frames, :64]
|
||||
|
||||
indiv_mels = torch.FloatTensor(indiv_mels).unsqueeze(1).unsqueeze(0) # bs T 1 80 16
|
||||
|
||||
if use_blink:
|
||||
ratio = torch.FloatTensor(ratio).unsqueeze(0) # bs T
|
||||
else:
|
||||
ratio = torch.FloatTensor(ratio).unsqueeze(0).fill_(0.)
|
||||
# bs T
|
||||
ref_coeff = torch.FloatTensor(ref_coeff).unsqueeze(0) # bs 1 70
|
||||
|
||||
indiv_mels = indiv_mels.to(device)
|
||||
ratio = ratio.to(device)
|
||||
ref_coeff = ref_coeff.to(device)
|
||||
|
||||
return {'indiv_mels': indiv_mels,
|
||||
'ref': ref_coeff,
|
||||
'num_frames': num_frames,
|
||||
'ratio_gt': ratio,
|
||||
'audio_name': audio_name, 'pic_name': pic_name}
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
import os
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from skimage import io, img_as_float32, transform
|
||||
import torch
|
||||
import scipy.io as scio
|
||||
|
||||
def get_facerender_data(coeff_path, pic_path, first_coeff_path, audio_path,
|
||||
batch_size, input_yaw_list=None, input_pitch_list=None, input_roll_list=None,
|
||||
expression_scale=1.0, still_mode = False, preprocess='crop', size = 256):
|
||||
|
||||
semantic_radius = 13
|
||||
video_name = os.path.splitext(os.path.split(coeff_path)[-1])[0]
|
||||
txt_path = os.path.splitext(coeff_path)[0]
|
||||
|
||||
data={}
|
||||
|
||||
img1 = Image.open(pic_path)
|
||||
source_image = np.array(img1)
|
||||
source_image = img_as_float32(source_image)
|
||||
source_image = transform.resize(source_image, (size, size, 3))
|
||||
source_image = source_image.transpose((2, 0, 1))
|
||||
source_image_ts = torch.FloatTensor(source_image).unsqueeze(0)
|
||||
source_image_ts = source_image_ts.repeat(batch_size, 1, 1, 1)
|
||||
data['source_image'] = source_image_ts
|
||||
|
||||
source_semantics_dict = scio.loadmat(first_coeff_path)
|
||||
generated_dict = scio.loadmat(coeff_path)
|
||||
|
||||
if 'full' not in preprocess.lower():
|
||||
source_semantics = source_semantics_dict['coeff_3dmm'][:1,:70] #1 70
|
||||
generated_3dmm = generated_dict['coeff_3dmm'][:,:70]
|
||||
|
||||
else:
|
||||
source_semantics = source_semantics_dict['coeff_3dmm'][:1,:73] #1 70
|
||||
generated_3dmm = generated_dict['coeff_3dmm'][:,:70]
|
||||
|
||||
source_semantics_new = transform_semantic_1(source_semantics, semantic_radius)
|
||||
source_semantics_ts = torch.FloatTensor(source_semantics_new).unsqueeze(0)
|
||||
source_semantics_ts = source_semantics_ts.repeat(batch_size, 1, 1)
|
||||
data['source_semantics'] = source_semantics_ts
|
||||
|
||||
# target
|
||||
generated_3dmm[:, :64] = generated_3dmm[:, :64] * expression_scale
|
||||
|
||||
if 'full' in preprocess.lower():
|
||||
generated_3dmm = np.concatenate([generated_3dmm, np.repeat(source_semantics[:,70:], generated_3dmm.shape[0], axis=0)], axis=1)
|
||||
|
||||
if still_mode:
|
||||
generated_3dmm[:, 64:] = np.repeat(source_semantics[:, 64:], generated_3dmm.shape[0], axis=0)
|
||||
|
||||
with open(txt_path+'.txt', 'w') as f:
|
||||
for coeff in generated_3dmm:
|
||||
for i in coeff:
|
||||
f.write(str(i)[:7] + ' '+'\t')
|
||||
f.write('\n')
|
||||
|
||||
target_semantics_list = []
|
||||
frame_num = generated_3dmm.shape[0]
|
||||
data['frame_num'] = frame_num
|
||||
for frame_idx in range(frame_num):
|
||||
target_semantics = transform_semantic_target(generated_3dmm, frame_idx, semantic_radius)
|
||||
target_semantics_list.append(target_semantics)
|
||||
|
||||
remainder = frame_num%batch_size
|
||||
if remainder!=0:
|
||||
for _ in range(batch_size-remainder):
|
||||
target_semantics_list.append(target_semantics)
|
||||
|
||||
target_semantics_np = np.array(target_semantics_list) #frame_num 70 semantic_radius*2+1
|
||||
target_semantics_np = target_semantics_np.reshape(batch_size, -1, target_semantics_np.shape[-2], target_semantics_np.shape[-1])
|
||||
data['target_semantics_list'] = torch.FloatTensor(target_semantics_np)
|
||||
data['video_name'] = video_name
|
||||
data['audio_path'] = audio_path
|
||||
|
||||
if input_yaw_list is not None:
|
||||
yaw_c_seq = gen_camera_pose(input_yaw_list, frame_num, batch_size)
|
||||
data['yaw_c_seq'] = torch.FloatTensor(yaw_c_seq)
|
||||
if input_pitch_list is not None:
|
||||
pitch_c_seq = gen_camera_pose(input_pitch_list, frame_num, batch_size)
|
||||
data['pitch_c_seq'] = torch.FloatTensor(pitch_c_seq)
|
||||
if input_roll_list is not None:
|
||||
roll_c_seq = gen_camera_pose(input_roll_list, frame_num, batch_size)
|
||||
data['roll_c_seq'] = torch.FloatTensor(roll_c_seq)
|
||||
|
||||
return data
|
||||
|
||||
def transform_semantic_1(semantic, semantic_radius):
|
||||
semantic_list = [semantic for i in range(0, semantic_radius*2+1)]
|
||||
coeff_3dmm = np.concatenate(semantic_list, 0)
|
||||
return coeff_3dmm.transpose(1,0)
|
||||
|
||||
def transform_semantic_target(coeff_3dmm, frame_index, semantic_radius):
|
||||
num_frames = coeff_3dmm.shape[0]
|
||||
seq = list(range(frame_index- semantic_radius, frame_index + semantic_radius+1))
|
||||
index = [ min(max(item, 0), num_frames-1) for item in seq ]
|
||||
coeff_3dmm_g = coeff_3dmm[index, :]
|
||||
return coeff_3dmm_g.transpose(1,0)
|
||||
|
||||
def gen_camera_pose(camera_degree_list, frame_num, batch_size):
|
||||
|
||||
new_degree_list = []
|
||||
if len(camera_degree_list) == 1:
|
||||
for _ in range(frame_num):
|
||||
new_degree_list.append(camera_degree_list[0])
|
||||
remainder = frame_num%batch_size
|
||||
if remainder!=0:
|
||||
for _ in range(batch_size-remainder):
|
||||
new_degree_list.append(new_degree_list[-1])
|
||||
new_degree_np = np.array(new_degree_list).reshape(batch_size, -1)
|
||||
return new_degree_np
|
||||
|
||||
degree_sum = 0.
|
||||
for i, degree in enumerate(camera_degree_list[1:]):
|
||||
degree_sum += abs(degree-camera_degree_list[i])
|
||||
|
||||
degree_per_frame = degree_sum/(frame_num-1)
|
||||
for i, degree in enumerate(camera_degree_list[1:]):
|
||||
degree_last = camera_degree_list[i]
|
||||
degree_step = degree_per_frame * abs(degree-degree_last)/(degree-degree_last)
|
||||
new_degree_list = new_degree_list + list(np.arange(degree_last, degree, degree_step))
|
||||
if len(new_degree_list) > frame_num:
|
||||
new_degree_list = new_degree_list[:frame_num]
|
||||
elif len(new_degree_list) < frame_num:
|
||||
for _ in range(frame_num-len(new_degree_list)):
|
||||
new_degree_list.append(new_degree_list[-1])
|
||||
print(len(new_degree_list))
|
||||
print(frame_num)
|
||||
|
||||
remainder = frame_num%batch_size
|
||||
if remainder!=0:
|
||||
for _ in range(batch_size-remainder):
|
||||
new_degree_list.append(new_degree_list[-1])
|
||||
new_degree_np = np.array(new_degree_list).reshape(batch_size, -1)
|
||||
return new_degree_np
|
||||
|
||||
@@ -0,0 +1,155 @@
|
||||
import torch, uuid
|
||||
import os, sys, shutil
|
||||
from src.utils.preprocess import CropAndExtract
|
||||
from src.test_audio2coeff import Audio2Coeff
|
||||
from src.facerender.animate import AnimateFromCoeff
|
||||
from src.generate_batch import get_data
|
||||
from src.generate_facerender_batch import get_facerender_data
|
||||
|
||||
from src.utils.init_path import init_path
|
||||
|
||||
from pydub import AudioSegment
|
||||
|
||||
|
||||
def mp3_to_wav(mp3_filename,wav_filename,frame_rate):
|
||||
mp3_file = AudioSegment.from_file(file=mp3_filename)
|
||||
mp3_file.set_frame_rate(frame_rate).export(wav_filename,format="wav")
|
||||
|
||||
|
||||
class SadTalker():
|
||||
|
||||
def __init__(self, checkpoint_path='checkpoints', config_path='src/config', lazy_load=False):
|
||||
|
||||
if torch.cuda.is_available() :
|
||||
device = "cuda"
|
||||
else:
|
||||
device = "cpu"
|
||||
|
||||
self.device = device
|
||||
|
||||
os.environ['TORCH_HOME']= checkpoint_path
|
||||
|
||||
self.checkpoint_path = checkpoint_path
|
||||
self.config_path = config_path
|
||||
|
||||
|
||||
def test(self, source_image, driven_audio, preprocess='crop',
|
||||
still_mode=False, use_enhancer=False, batch_size=1, size=256,
|
||||
pose_style = 0, exp_scale=1.0,
|
||||
use_ref_video = False,
|
||||
ref_video = None,
|
||||
ref_info = None,
|
||||
use_idle_mode = False,
|
||||
length_of_audio = 0, use_blink=True,
|
||||
result_dir='./results/'):
|
||||
|
||||
self.sadtalker_paths = init_path(self.checkpoint_path, self.config_path, size, False, preprocess)
|
||||
print(self.sadtalker_paths)
|
||||
|
||||
self.audio_to_coeff = Audio2Coeff(self.sadtalker_paths, self.device)
|
||||
self.preprocess_model = CropAndExtract(self.sadtalker_paths, self.device)
|
||||
self.animate_from_coeff = AnimateFromCoeff(self.sadtalker_paths, self.device)
|
||||
|
||||
time_tag = str(uuid.uuid4())
|
||||
save_dir = os.path.join(result_dir, time_tag)
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
input_dir = os.path.join(save_dir, 'input')
|
||||
os.makedirs(input_dir, exist_ok=True)
|
||||
|
||||
print(source_image)
|
||||
pic_path = os.path.join(input_dir, os.path.basename(source_image))
|
||||
shutil.move(source_image, input_dir)
|
||||
|
||||
if driven_audio is not None and os.path.isfile(driven_audio):
|
||||
audio_path = os.path.join(input_dir, os.path.basename(driven_audio))
|
||||
|
||||
#### mp3 to wav
|
||||
if '.mp3' in audio_path:
|
||||
mp3_to_wav(driven_audio, audio_path.replace('.mp3', '.wav'), 16000)
|
||||
audio_path = audio_path.replace('.mp3', '.wav')
|
||||
else:
|
||||
shutil.move(driven_audio, input_dir)
|
||||
|
||||
elif use_idle_mode:
|
||||
audio_path = os.path.join(input_dir, 'idlemode_'+str(length_of_audio)+'.wav') ## generate audio from this new audio_path
|
||||
from pydub import AudioSegment
|
||||
one_sec_segment = AudioSegment.silent(duration=1000*length_of_audio) #duration in milliseconds
|
||||
one_sec_segment.export(audio_path, format="wav")
|
||||
else:
|
||||
print(use_ref_video, ref_info)
|
||||
assert use_ref_video == True and ref_info == 'all'
|
||||
|
||||
if use_ref_video and ref_info == 'all': # full ref mode
|
||||
ref_video_videoname = os.path.basename(ref_video)
|
||||
audio_path = os.path.join(save_dir, ref_video_videoname+'.wav')
|
||||
print('new audiopath:',audio_path)
|
||||
# if ref_video contains audio, set the audio from ref_video.
|
||||
cmd = r"ffmpeg -y -hide_banner -loglevel error -i %s %s"%(ref_video, audio_path)
|
||||
os.system(cmd)
|
||||
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
#crop image and extract 3dmm from image
|
||||
first_frame_dir = os.path.join(save_dir, 'first_frame_dir')
|
||||
os.makedirs(first_frame_dir, exist_ok=True)
|
||||
first_coeff_path, crop_pic_path, crop_info = self.preprocess_model.generate(pic_path, first_frame_dir, preprocess, True, size)
|
||||
|
||||
if first_coeff_path is None:
|
||||
raise AttributeError("No face is detected")
|
||||
|
||||
if use_ref_video:
|
||||
print('using ref video for genreation')
|
||||
ref_video_videoname = os.path.splitext(os.path.split(ref_video)[-1])[0]
|
||||
ref_video_frame_dir = os.path.join(save_dir, ref_video_videoname)
|
||||
os.makedirs(ref_video_frame_dir, exist_ok=True)
|
||||
print('3DMM Extraction for the reference video providing pose')
|
||||
ref_video_coeff_path, _, _ = self.preprocess_model.generate(ref_video, ref_video_frame_dir, preprocess, source_image_flag=False)
|
||||
else:
|
||||
ref_video_coeff_path = None
|
||||
|
||||
if use_ref_video:
|
||||
if ref_info == 'pose':
|
||||
ref_pose_coeff_path = ref_video_coeff_path
|
||||
ref_eyeblink_coeff_path = None
|
||||
elif ref_info == 'blink':
|
||||
ref_pose_coeff_path = None
|
||||
ref_eyeblink_coeff_path = ref_video_coeff_path
|
||||
elif ref_info == 'pose+blink':
|
||||
ref_pose_coeff_path = ref_video_coeff_path
|
||||
ref_eyeblink_coeff_path = ref_video_coeff_path
|
||||
elif ref_info == 'all':
|
||||
ref_pose_coeff_path = None
|
||||
ref_eyeblink_coeff_path = None
|
||||
else:
|
||||
raise('error in refinfo')
|
||||
else:
|
||||
ref_pose_coeff_path = None
|
||||
ref_eyeblink_coeff_path = None
|
||||
|
||||
#audio2ceoff
|
||||
if use_ref_video and ref_info == 'all':
|
||||
coeff_path = ref_video_coeff_path # self.audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path)
|
||||
else:
|
||||
batch = get_data(first_coeff_path, audio_path, self.device, ref_eyeblink_coeff_path=ref_eyeblink_coeff_path, still=still_mode, idlemode=use_idle_mode, length_of_audio=length_of_audio, use_blink=use_blink) # longer audio?
|
||||
coeff_path = self.audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path)
|
||||
|
||||
#coeff2video
|
||||
data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, batch_size, still_mode=still_mode, preprocess=preprocess, size=size, expression_scale = exp_scale)
|
||||
return_path = self.animate_from_coeff.generate(data, save_dir, pic_path, crop_info, enhancer='gfpgan' if use_enhancer else None, preprocess=preprocess, img_size=size)
|
||||
video_name = data['video_name']
|
||||
print(f'The generated video is named {video_name} in {save_dir}')
|
||||
|
||||
del self.preprocess_model
|
||||
del self.audio_to_coeff
|
||||
del self.animate_from_coeff
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
import gc; gc.collect()
|
||||
|
||||
return return_path
|
||||
|
||||
|
||||
@@ -0,0 +1,123 @@
|
||||
import os
|
||||
import torch
|
||||
import numpy as np
|
||||
from scipy.io import savemat, loadmat
|
||||
from yacs.config import CfgNode as CN
|
||||
from scipy.signal import savgol_filter
|
||||
|
||||
import safetensors
|
||||
import safetensors.torch
|
||||
|
||||
from src.audio2pose_models.audio2pose import Audio2Pose
|
||||
from src.audio2exp_models.networks import SimpleWrapperV2
|
||||
from src.audio2exp_models.audio2exp import Audio2Exp
|
||||
from src.utils.safetensor_helper import load_x_from_safetensor
|
||||
|
||||
def load_cpk(checkpoint_path, model=None, optimizer=None, device="cpu"):
|
||||
checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
|
||||
if model is not None:
|
||||
model.load_state_dict(checkpoint['model'])
|
||||
if optimizer is not None:
|
||||
optimizer.load_state_dict(checkpoint['optimizer'])
|
||||
|
||||
return checkpoint['epoch']
|
||||
|
||||
class Audio2Coeff():
|
||||
|
||||
def __init__(self, sadtalker_path, device):
|
||||
#load config
|
||||
fcfg_pose = open(sadtalker_path['audio2pose_yaml_path'])
|
||||
cfg_pose = CN.load_cfg(fcfg_pose)
|
||||
cfg_pose.freeze()
|
||||
fcfg_exp = open(sadtalker_path['audio2exp_yaml_path'])
|
||||
cfg_exp = CN.load_cfg(fcfg_exp)
|
||||
cfg_exp.freeze()
|
||||
|
||||
# load audio2pose_model
|
||||
self.audio2pose_model = Audio2Pose(cfg_pose, None, device=device)
|
||||
self.audio2pose_model = self.audio2pose_model.to(device)
|
||||
self.audio2pose_model.eval()
|
||||
for param in self.audio2pose_model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
try:
|
||||
if sadtalker_path['use_safetensor']:
|
||||
checkpoints = safetensors.torch.load_file(sadtalker_path['checkpoint'])
|
||||
self.audio2pose_model.load_state_dict(load_x_from_safetensor(checkpoints, 'audio2pose'))
|
||||
else:
|
||||
load_cpk(sadtalker_path['audio2pose_checkpoint'], model=self.audio2pose_model, device=device)
|
||||
except:
|
||||
raise Exception("Failed in loading audio2pose_checkpoint")
|
||||
|
||||
# load audio2exp_model
|
||||
netG = SimpleWrapperV2()
|
||||
netG = netG.to(device)
|
||||
for param in netG.parameters():
|
||||
netG.requires_grad = False
|
||||
netG.eval()
|
||||
try:
|
||||
if sadtalker_path['use_safetensor']:
|
||||
checkpoints = safetensors.torch.load_file(sadtalker_path['checkpoint'])
|
||||
netG.load_state_dict(load_x_from_safetensor(checkpoints, 'audio2exp'))
|
||||
else:
|
||||
load_cpk(sadtalker_path['audio2exp_checkpoint'], model=netG, device=device)
|
||||
except:
|
||||
raise Exception("Failed in loading audio2exp_checkpoint")
|
||||
self.audio2exp_model = Audio2Exp(netG, cfg_exp, device=device, prepare_training_loss=False)
|
||||
self.audio2exp_model = self.audio2exp_model.to(device)
|
||||
for param in self.audio2exp_model.parameters():
|
||||
param.requires_grad = False
|
||||
self.audio2exp_model.eval()
|
||||
|
||||
self.device = device
|
||||
|
||||
def generate(self, batch, coeff_save_dir, pose_style, ref_pose_coeff_path=None):
|
||||
|
||||
with torch.no_grad():
|
||||
#test
|
||||
results_dict_exp= self.audio2exp_model.test(batch)
|
||||
exp_pred = results_dict_exp['exp_coeff_pred'] #bs T 64
|
||||
|
||||
#for class_id in range(1):
|
||||
#class_id = 0#(i+10)%45
|
||||
#class_id = random.randint(0,46) #46 styles can be selected
|
||||
batch['class'] = torch.LongTensor([pose_style]).to(self.device)
|
||||
results_dict_pose = self.audio2pose_model.test(batch)
|
||||
pose_pred = results_dict_pose['pose_pred'] #bs T 6
|
||||
|
||||
pose_len = pose_pred.shape[1]
|
||||
if pose_len<13:
|
||||
pose_len = int((pose_len-1)/2)*2+1
|
||||
pose_pred = torch.Tensor(savgol_filter(np.array(pose_pred.cpu()), pose_len, 2, axis=1)).to(self.device)
|
||||
else:
|
||||
pose_pred = torch.Tensor(savgol_filter(np.array(pose_pred.cpu()), 13, 2, axis=1)).to(self.device)
|
||||
|
||||
coeffs_pred = torch.cat((exp_pred, pose_pred), dim=-1) #bs T 70
|
||||
|
||||
coeffs_pred_numpy = coeffs_pred[0].clone().detach().cpu().numpy()
|
||||
|
||||
if ref_pose_coeff_path is not None:
|
||||
coeffs_pred_numpy = self.using_refpose(coeffs_pred_numpy, ref_pose_coeff_path)
|
||||
|
||||
savemat(os.path.join(coeff_save_dir, '%s##%s.mat'%(batch['pic_name'], batch['audio_name'])),
|
||||
{'coeff_3dmm': coeffs_pred_numpy})
|
||||
|
||||
return os.path.join(coeff_save_dir, '%s##%s.mat'%(batch['pic_name'], batch['audio_name']))
|
||||
|
||||
def using_refpose(self, coeffs_pred_numpy, ref_pose_coeff_path):
|
||||
num_frames = coeffs_pred_numpy.shape[0]
|
||||
refpose_coeff_dict = loadmat(ref_pose_coeff_path)
|
||||
refpose_coeff = refpose_coeff_dict['coeff_3dmm'][:,64:70]
|
||||
refpose_num_frames = refpose_coeff.shape[0]
|
||||
if refpose_num_frames<num_frames:
|
||||
div = num_frames//refpose_num_frames
|
||||
re = num_frames%refpose_num_frames
|
||||
refpose_coeff_list = [refpose_coeff for i in range(div)]
|
||||
refpose_coeff_list.append(refpose_coeff[:re, :])
|
||||
refpose_coeff = np.concatenate(refpose_coeff_list, axis=0)
|
||||
|
||||
#### relative head pose
|
||||
coeffs_pred_numpy[:, 64:70] = coeffs_pred_numpy[:, 64:70] + ( refpose_coeff[:num_frames, :] - refpose_coeff[0:1, :] )
|
||||
return coeffs_pred_numpy
|
||||
|
||||
|
||||
Reference in New Issue
Block a user