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
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,136 @@
|
||||
import librosa
|
||||
import librosa.filters
|
||||
import numpy as np
|
||||
# import tensorflow as tf
|
||||
from scipy import signal
|
||||
from scipy.io import wavfile
|
||||
from src.utils.hparams import hparams as hp
|
||||
|
||||
def load_wav(path, sr):
|
||||
return librosa.core.load(path, sr=sr)[0]
|
||||
|
||||
def save_wav(wav, path, sr):
|
||||
wav *= 32767 / max(0.01, np.max(np.abs(wav)))
|
||||
#proposed by @dsmiller
|
||||
wavfile.write(path, sr, wav.astype(np.int16))
|
||||
|
||||
def save_wavenet_wav(wav, path, sr):
|
||||
librosa.output.write_wav(path, wav, sr=sr)
|
||||
|
||||
def preemphasis(wav, k, preemphasize=True):
|
||||
if preemphasize:
|
||||
return signal.lfilter([1, -k], [1], wav)
|
||||
return wav
|
||||
|
||||
def inv_preemphasis(wav, k, inv_preemphasize=True):
|
||||
if inv_preemphasize:
|
||||
return signal.lfilter([1], [1, -k], wav)
|
||||
return wav
|
||||
|
||||
def get_hop_size():
|
||||
hop_size = hp.hop_size
|
||||
if hop_size is None:
|
||||
assert hp.frame_shift_ms is not None
|
||||
hop_size = int(hp.frame_shift_ms / 1000 * hp.sample_rate)
|
||||
return hop_size
|
||||
|
||||
def linearspectrogram(wav):
|
||||
D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize))
|
||||
S = _amp_to_db(np.abs(D)) - hp.ref_level_db
|
||||
|
||||
if hp.signal_normalization:
|
||||
return _normalize(S)
|
||||
return S
|
||||
|
||||
def melspectrogram(wav):
|
||||
D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize))
|
||||
S = _amp_to_db(_linear_to_mel(np.abs(D))) - hp.ref_level_db
|
||||
|
||||
if hp.signal_normalization:
|
||||
return _normalize(S)
|
||||
return S
|
||||
|
||||
def _lws_processor():
|
||||
import lws
|
||||
return lws.lws(hp.n_fft, get_hop_size(), fftsize=hp.win_size, mode="speech")
|
||||
|
||||
def _stft(y):
|
||||
if hp.use_lws:
|
||||
return _lws_processor(hp).stft(y).T
|
||||
else:
|
||||
return librosa.stft(y=y, n_fft=hp.n_fft, hop_length=get_hop_size(), win_length=hp.win_size)
|
||||
|
||||
##########################################################
|
||||
#Those are only correct when using lws!!! (This was messing with Wavenet quality for a long time!)
|
||||
def num_frames(length, fsize, fshift):
|
||||
"""Compute number of time frames of spectrogram
|
||||
"""
|
||||
pad = (fsize - fshift)
|
||||
if length % fshift == 0:
|
||||
M = (length + pad * 2 - fsize) // fshift + 1
|
||||
else:
|
||||
M = (length + pad * 2 - fsize) // fshift + 2
|
||||
return M
|
||||
|
||||
|
||||
def pad_lr(x, fsize, fshift):
|
||||
"""Compute left and right padding
|
||||
"""
|
||||
M = num_frames(len(x), fsize, fshift)
|
||||
pad = (fsize - fshift)
|
||||
T = len(x) + 2 * pad
|
||||
r = (M - 1) * fshift + fsize - T
|
||||
return pad, pad + r
|
||||
##########################################################
|
||||
#Librosa correct padding
|
||||
def librosa_pad_lr(x, fsize, fshift):
|
||||
return 0, (x.shape[0] // fshift + 1) * fshift - x.shape[0]
|
||||
|
||||
# Conversions
|
||||
_mel_basis = None
|
||||
|
||||
def _linear_to_mel(spectogram):
|
||||
global _mel_basis
|
||||
if _mel_basis is None:
|
||||
_mel_basis = _build_mel_basis()
|
||||
return np.dot(_mel_basis, spectogram)
|
||||
|
||||
def _build_mel_basis():
|
||||
assert hp.fmax <= hp.sample_rate // 2
|
||||
return librosa.filters.mel(sr=hp.sample_rate, n_fft=hp.n_fft, n_mels=hp.num_mels,
|
||||
fmin=hp.fmin, fmax=hp.fmax)
|
||||
|
||||
def _amp_to_db(x):
|
||||
min_level = np.exp(hp.min_level_db / 20 * np.log(10))
|
||||
return 20 * np.log10(np.maximum(min_level, x))
|
||||
|
||||
def _db_to_amp(x):
|
||||
return np.power(10.0, (x) * 0.05)
|
||||
|
||||
def _normalize(S):
|
||||
if hp.allow_clipping_in_normalization:
|
||||
if hp.symmetric_mels:
|
||||
return np.clip((2 * hp.max_abs_value) * ((S - hp.min_level_db) / (-hp.min_level_db)) - hp.max_abs_value,
|
||||
-hp.max_abs_value, hp.max_abs_value)
|
||||
else:
|
||||
return np.clip(hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db)), 0, hp.max_abs_value)
|
||||
|
||||
assert S.max() <= 0 and S.min() - hp.min_level_db >= 0
|
||||
if hp.symmetric_mels:
|
||||
return (2 * hp.max_abs_value) * ((S - hp.min_level_db) / (-hp.min_level_db)) - hp.max_abs_value
|
||||
else:
|
||||
return hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db))
|
||||
|
||||
def _denormalize(D):
|
||||
if hp.allow_clipping_in_normalization:
|
||||
if hp.symmetric_mels:
|
||||
return (((np.clip(D, -hp.max_abs_value,
|
||||
hp.max_abs_value) + hp.max_abs_value) * -hp.min_level_db / (2 * hp.max_abs_value))
|
||||
+ hp.min_level_db)
|
||||
else:
|
||||
return ((np.clip(D, 0, hp.max_abs_value) * -hp.min_level_db / hp.max_abs_value) + hp.min_level_db)
|
||||
|
||||
if hp.symmetric_mels:
|
||||
return (((D + hp.max_abs_value) * -hp.min_level_db / (2 * hp.max_abs_value)) + hp.min_level_db)
|
||||
else:
|
||||
return ((D * -hp.min_level_db / hp.max_abs_value) + hp.min_level_db)
|
||||
@@ -0,0 +1,144 @@
|
||||
import os
|
||||
import cv2
|
||||
import time
|
||||
import glob
|
||||
import argparse
|
||||
import scipy
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
from itertools import cycle
|
||||
|
||||
from src.face3d.extract_kp_videos_safe import KeypointExtractor
|
||||
from facexlib.alignment import landmark_98_to_68
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
class Preprocesser:
|
||||
def __init__(self, device='cuda'):
|
||||
self.predictor = KeypointExtractor(device)
|
||||
|
||||
def get_landmark(self, img_np):
|
||||
"""get landmark with dlib
|
||||
:return: np.array shape=(68, 2)
|
||||
"""
|
||||
with torch.no_grad():
|
||||
dets = self.predictor.det_net.detect_faces(img_np, 0.97)
|
||||
|
||||
if len(dets) == 0:
|
||||
return None
|
||||
det = dets[0]
|
||||
|
||||
img = img_np[int(det[1]):int(det[3]), int(det[0]):int(det[2]), :]
|
||||
lm = landmark_98_to_68(self.predictor.detector.get_landmarks(img)) # [0]
|
||||
|
||||
#### keypoints to the original location
|
||||
lm[:,0] += int(det[0])
|
||||
lm[:,1] += int(det[1])
|
||||
|
||||
return lm
|
||||
|
||||
def align_face(self, img, lm, output_size=1024):
|
||||
"""
|
||||
:param filepath: str
|
||||
:return: PIL Image
|
||||
"""
|
||||
lm_chin = lm[0: 17] # left-right
|
||||
lm_eyebrow_left = lm[17: 22] # left-right
|
||||
lm_eyebrow_right = lm[22: 27] # left-right
|
||||
lm_nose = lm[27: 31] # top-down
|
||||
lm_nostrils = lm[31: 36] # top-down
|
||||
lm_eye_left = lm[36: 42] # left-clockwise
|
||||
lm_eye_right = lm[42: 48] # left-clockwise
|
||||
lm_mouth_outer = lm[48: 60] # left-clockwise
|
||||
lm_mouth_inner = lm[60: 68] # left-clockwise
|
||||
|
||||
# Calculate auxiliary vectors.
|
||||
eye_left = np.mean(lm_eye_left, axis=0)
|
||||
eye_right = np.mean(lm_eye_right, axis=0)
|
||||
eye_avg = (eye_left + eye_right) * 0.5
|
||||
eye_to_eye = eye_right - eye_left
|
||||
mouth_left = lm_mouth_outer[0]
|
||||
mouth_right = lm_mouth_outer[6]
|
||||
mouth_avg = (mouth_left + mouth_right) * 0.5
|
||||
eye_to_mouth = mouth_avg - eye_avg
|
||||
|
||||
# Choose oriented crop rectangle.
|
||||
x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1] # Addition of binocular difference and double mouth difference
|
||||
x /= np.hypot(*x) # hypot函数计算直角三角形的斜边长，用斜边长对三角形两条直边做归一化
|
||||
x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8) # 双眼差和眼嘴差，选较大的作为基准尺度
|
||||
y = np.flipud(x) * [-1, 1]
|
||||
c = eye_avg + eye_to_mouth * 0.1
|
||||
quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y]) # 定义四边形，以面部基准位置为中心上下左右平移得到四个顶点
|
||||
qsize = np.hypot(*x) * 2 # 定义四边形的大小（边长），为基准尺度的2倍
|
||||
|
||||
# Shrink.
|
||||
# 如果计算出的四边形太大了，就按比例缩小它
|
||||
shrink = int(np.floor(qsize / output_size * 0.5))
|
||||
if shrink > 1:
|
||||
rsize = (int(np.rint(float(img.size[0]) / shrink)), int(np.rint(float(img.size[1]) / shrink)))
|
||||
img = img.resize(rsize, Image.ANTIALIAS)
|
||||
quad /= shrink
|
||||
qsize /= shrink
|
||||
else:
|
||||
rsize = (int(np.rint(float(img.size[0]))), int(np.rint(float(img.size[1]))))
|
||||
|
||||
# Crop.
|
||||
border = max(int(np.rint(qsize * 0.1)), 3)
|
||||
crop = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
|
||||
int(np.ceil(max(quad[:, 1]))))
|
||||
crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, img.size[0]),
|
||||
min(crop[3] + border, img.size[1]))
|
||||
if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]:
|
||||
# img = img.crop(crop)
|
||||
quad -= crop[0:2]
|
||||
|
||||
# Pad.
|
||||
pad = (int(np.floor(min(quad[:, 0]))), int(np.floor(min(quad[:, 1]))), int(np.ceil(max(quad[:, 0]))),
|
||||
int(np.ceil(max(quad[:, 1]))))
|
||||
pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - img.size[0] + border, 0),
|
||||
max(pad[3] - img.size[1] + border, 0))
|
||||
# if enable_padding and max(pad) > border - 4:
|
||||
# pad = np.maximum(pad, int(np.rint(qsize * 0.3)))
|
||||
# img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')
|
||||
# h, w, _ = img.shape
|
||||
# y, x, _ = np.ogrid[:h, :w, :1]
|
||||
# mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w - 1 - x) / pad[2]),
|
||||
# 1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h - 1 - y) / pad[3]))
|
||||
# blur = qsize * 0.02
|
||||
# img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
|
||||
# img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0)
|
||||
# img = Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), 'RGB')
|
||||
# quad += pad[:2]
|
||||
|
||||
# Transform.
|
||||
quad = (quad + 0.5).flatten()
|
||||
lx = max(min(quad[0], quad[2]), 0)
|
||||
ly = max(min(quad[1], quad[7]), 0)
|
||||
rx = min(max(quad[4], quad[6]), img.size[0])
|
||||
ry = min(max(quad[3], quad[5]), img.size[0])
|
||||
|
||||
# Save aligned image.
|
||||
return rsize, crop, [lx, ly, rx, ry]
|
||||
|
||||
def crop(self, img_np_list, still=False, xsize=512): # first frame for all video
|
||||
img_np = img_np_list[0]
|
||||
lm = self.get_landmark(img_np)
|
||||
|
||||
if lm is None:
|
||||
raise 'can not detect the landmark from source image'
|
||||
rsize, crop, quad = self.align_face(img=Image.fromarray(img_np), lm=lm, output_size=xsize)
|
||||
clx, cly, crx, cry = crop
|
||||
lx, ly, rx, ry = quad
|
||||
lx, ly, rx, ry = int(lx), int(ly), int(rx), int(ry)
|
||||
for _i in range(len(img_np_list)):
|
||||
_inp = img_np_list[_i]
|
||||
_inp = cv2.resize(_inp, (rsize[0], rsize[1]))
|
||||
_inp = _inp[cly:cry, clx:crx]
|
||||
if not still:
|
||||
_inp = _inp[ly:ry, lx:rx]
|
||||
img_np_list[_i] = _inp
|
||||
return img_np_list, crop, quad
|
||||
|
||||
@@ -0,0 +1,123 @@
|
||||
import os
|
||||
import torch
|
||||
|
||||
from gfpgan import GFPGANer
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from src.utils.videoio import load_video_to_cv2
|
||||
|
||||
import cv2
|
||||
|
||||
|
||||
class GeneratorWithLen(object):
|
||||
""" From https://stackoverflow.com/a/7460929 """
|
||||
|
||||
def __init__(self, gen, length):
|
||||
self.gen = gen
|
||||
self.length = length
|
||||
|
||||
def __len__(self):
|
||||
return self.length
|
||||
|
||||
def __iter__(self):
|
||||
return self.gen
|
||||
|
||||
def enhancer_list(images, method='gfpgan', bg_upsampler='realesrgan'):
|
||||
gen = enhancer_generator_no_len(images, method=method, bg_upsampler=bg_upsampler)
|
||||
return list(gen)
|
||||
|
||||
def enhancer_generator_with_len(images, method='gfpgan', bg_upsampler='realesrgan'):
|
||||
""" Provide a generator with a __len__ method so that it can passed to functions that
|
||||
call len()"""
|
||||
|
||||
if os.path.isfile(images): # handle video to images
|
||||
# TODO: Create a generator version of load_video_to_cv2
|
||||
images = load_video_to_cv2(images)
|
||||
|
||||
gen = enhancer_generator_no_len(images, method=method, bg_upsampler=bg_upsampler)
|
||||
gen_with_len = GeneratorWithLen(gen, len(images))
|
||||
return gen_with_len
|
||||
|
||||
def enhancer_generator_no_len(images, method='gfpgan', bg_upsampler='realesrgan'):
|
||||
""" Provide a generator function so that all of the enhanced images don't need
|
||||
to be stored in memory at the same time. This can save tons of RAM compared to
|
||||
the enhancer function. """
|
||||
|
||||
print('face enhancer....')
|
||||
if not isinstance(images, list) and os.path.isfile(images): # handle video to images
|
||||
images = load_video_to_cv2(images)
|
||||
|
||||
# ------------------------ set up GFPGAN restorer ------------------------
|
||||
if method == 'gfpgan':
|
||||
arch = 'clean'
|
||||
channel_multiplier = 2
|
||||
model_name = 'GFPGANv1.4'
|
||||
url = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth'
|
||||
elif method == 'RestoreFormer':
|
||||
arch = 'RestoreFormer'
|
||||
channel_multiplier = 2
|
||||
model_name = 'RestoreFormer'
|
||||
url = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth'
|
||||
elif method == 'codeformer': # TODO:
|
||||
arch = 'CodeFormer'
|
||||
channel_multiplier = 2
|
||||
model_name = 'CodeFormer'
|
||||
url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
|
||||
else:
|
||||
raise ValueError(f'Wrong model version {method}.')
|
||||
|
||||
|
||||
# ------------------------ set up background upsampler ------------------------
|
||||
if bg_upsampler == 'realesrgan':
|
||||
if not torch.cuda.is_available(): # CPU
|
||||
import warnings
|
||||
warnings.warn('The unoptimized RealESRGAN is slow on CPU. We do not use it. '
|
||||
'If you really want to use it, please modify the corresponding codes.')
|
||||
bg_upsampler = None
|
||||
else:
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet
|
||||
from realesrgan import RealESRGANer
|
||||
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
|
||||
bg_upsampler = RealESRGANer(
|
||||
scale=2,
|
||||
model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth',
|
||||
model=model,
|
||||
tile=400,
|
||||
tile_pad=10,
|
||||
pre_pad=0,
|
||||
half=True) # need to set False in CPU mode
|
||||
else:
|
||||
bg_upsampler = None
|
||||
|
||||
# determine model paths
|
||||
model_path = os.path.join('gfpgan/weights', model_name + '.pth')
|
||||
|
||||
if not os.path.isfile(model_path):
|
||||
model_path = os.path.join('checkpoints', model_name + '.pth')
|
||||
|
||||
if not os.path.isfile(model_path):
|
||||
# download pre-trained models from url
|
||||
model_path = url
|
||||
|
||||
restorer = GFPGANer(
|
||||
model_path=model_path,
|
||||
upscale=2,
|
||||
arch=arch,
|
||||
channel_multiplier=channel_multiplier,
|
||||
bg_upsampler=bg_upsampler)
|
||||
|
||||
# ------------------------ restore ------------------------
|
||||
for idx in tqdm(range(len(images)), 'Face Enhancer:'):
|
||||
|
||||
img = cv2.cvtColor(images[idx], cv2.COLOR_RGB2BGR)
|
||||
|
||||
# restore faces and background if necessary
|
||||
cropped_faces, restored_faces, r_img = restorer.enhance(
|
||||
img,
|
||||
has_aligned=False,
|
||||
only_center_face=False,
|
||||
paste_back=True)
|
||||
|
||||
r_img = cv2.cvtColor(r_img, cv2.COLOR_BGR2RGB)
|
||||
yield r_img
|
||||
@@ -0,0 +1,160 @@
|
||||
from glob import glob
|
||||
import os
|
||||
|
||||
class HParams:
|
||||
def __init__(self, **kwargs):
|
||||
self.data = {}
|
||||
|
||||
for key, value in kwargs.items():
|
||||
self.data[key] = value
|
||||
|
||||
def __getattr__(self, key):
|
||||
if key not in self.data:
|
||||
raise AttributeError("'HParams' object has no attribute %s" % key)
|
||||
return self.data[key]
|
||||
|
||||
def set_hparam(self, key, value):
|
||||
self.data[key] = value
|
||||
|
||||
|
||||
# Default hyperparameters
|
||||
hparams = HParams(
|
||||
num_mels=80, # Number of mel-spectrogram channels and local conditioning dimensionality
|
||||
# network
|
||||
rescale=True, # Whether to rescale audio prior to preprocessing
|
||||
rescaling_max=0.9, # Rescaling value
|
||||
|
||||
# Use LWS (https://github.com/Jonathan-LeRoux/lws) for STFT and phase reconstruction
|
||||
# It"s preferred to set True to use with https://github.com/r9y9/wavenet_vocoder
|
||||
# Does not work if n_ffit is not multiple of hop_size!!
|
||||
use_lws=False,
|
||||
|
||||
n_fft=800, # Extra window size is filled with 0 paddings to match this parameter
|
||||
hop_size=200, # For 16000Hz, 200 = 12.5 ms (0.0125 * sample_rate)
|
||||
win_size=800, # For 16000Hz, 800 = 50 ms (If None, win_size = n_fft) (0.05 * sample_rate)
|
||||
sample_rate=16000, # 16000Hz (corresponding to librispeech) (sox --i <filename>)
|
||||
|
||||
frame_shift_ms=None, # Can replace hop_size parameter. (Recommended: 12.5)
|
||||
|
||||
# Mel and Linear spectrograms normalization/scaling and clipping
|
||||
signal_normalization=True,
|
||||
# Whether to normalize mel spectrograms to some predefined range (following below parameters)
|
||||
allow_clipping_in_normalization=True, # Only relevant if mel_normalization = True
|
||||
symmetric_mels=True,
|
||||
# Whether to scale the data to be symmetric around 0. (Also multiplies the output range by 2,
|
||||
# faster and cleaner convergence)
|
||||
max_abs_value=4.,
|
||||
# max absolute value of data. If symmetric, data will be [-max, max] else [0, max] (Must not
|
||||
# be too big to avoid gradient explosion,
|
||||
# not too small for fast convergence)
|
||||
# Contribution by @begeekmyfriend
|
||||
# Spectrogram Pre-Emphasis (Lfilter: Reduce spectrogram noise and helps model certitude
|
||||
# levels. Also allows for better G&L phase reconstruction)
|
||||
preemphasize=True, # whether to apply filter
|
||||
preemphasis=0.97, # filter coefficient.
|
||||
|
||||
# Limits
|
||||
min_level_db=-100,
|
||||
ref_level_db=20,
|
||||
fmin=55,
|
||||
# Set this to 55 if your speaker is male! if female, 95 should help taking off noise. (To
|
||||
# test depending on dataset. Pitch info: male~[65, 260], female~[100, 525])
|
||||
fmax=7600, # To be increased/reduced depending on data.
|
||||
|
||||
###################### Our training parameters #################################
|
||||
img_size=96,
|
||||
fps=25,
|
||||
|
||||
batch_size=16,
|
||||
initial_learning_rate=1e-4,
|
||||
nepochs=300000, ### ctrl + c, stop whenever eval loss is consistently greater than train loss for ~10 epochs
|
||||
num_workers=20,
|
||||
checkpoint_interval=3000,
|
||||
eval_interval=3000,
|
||||
writer_interval=300,
|
||||
save_optimizer_state=True,
|
||||
|
||||
syncnet_wt=0.0, # is initially zero, will be set automatically to 0.03 later. Leads to faster convergence.
|
||||
syncnet_batch_size=64,
|
||||
syncnet_lr=1e-4,
|
||||
syncnet_eval_interval=1000,
|
||||
syncnet_checkpoint_interval=10000,
|
||||
|
||||
disc_wt=0.07,
|
||||
disc_initial_learning_rate=1e-4,
|
||||
)
|
||||
|
||||
|
||||
|
||||
# Default hyperparameters
|
||||
hparamsdebug = HParams(
|
||||
num_mels=80, # Number of mel-spectrogram channels and local conditioning dimensionality
|
||||
# network
|
||||
rescale=True, # Whether to rescale audio prior to preprocessing
|
||||
rescaling_max=0.9, # Rescaling value
|
||||
|
||||
# Use LWS (https://github.com/Jonathan-LeRoux/lws) for STFT and phase reconstruction
|
||||
# It"s preferred to set True to use with https://github.com/r9y9/wavenet_vocoder
|
||||
# Does not work if n_ffit is not multiple of hop_size!!
|
||||
use_lws=False,
|
||||
|
||||
n_fft=800, # Extra window size is filled with 0 paddings to match this parameter
|
||||
hop_size=200, # For 16000Hz, 200 = 12.5 ms (0.0125 * sample_rate)
|
||||
win_size=800, # For 16000Hz, 800 = 50 ms (If None, win_size = n_fft) (0.05 * sample_rate)
|
||||
sample_rate=16000, # 16000Hz (corresponding to librispeech) (sox --i <filename>)
|
||||
|
||||
frame_shift_ms=None, # Can replace hop_size parameter. (Recommended: 12.5)
|
||||
|
||||
# Mel and Linear spectrograms normalization/scaling and clipping
|
||||
signal_normalization=True,
|
||||
# Whether to normalize mel spectrograms to some predefined range (following below parameters)
|
||||
allow_clipping_in_normalization=True, # Only relevant if mel_normalization = True
|
||||
symmetric_mels=True,
|
||||
# Whether to scale the data to be symmetric around 0. (Also multiplies the output range by 2,
|
||||
# faster and cleaner convergence)
|
||||
max_abs_value=4.,
|
||||
# max absolute value of data. If symmetric, data will be [-max, max] else [0, max] (Must not
|
||||
# be too big to avoid gradient explosion,
|
||||
# not too small for fast convergence)
|
||||
# Contribution by @begeekmyfriend
|
||||
# Spectrogram Pre-Emphasis (Lfilter: Reduce spectrogram noise and helps model certitude
|
||||
# levels. Also allows for better G&L phase reconstruction)
|
||||
preemphasize=True, # whether to apply filter
|
||||
preemphasis=0.97, # filter coefficient.
|
||||
|
||||
# Limits
|
||||
min_level_db=-100,
|
||||
ref_level_db=20,
|
||||
fmin=55,
|
||||
# Set this to 55 if your speaker is male! if female, 95 should help taking off noise. (To
|
||||
# test depending on dataset. Pitch info: male~[65, 260], female~[100, 525])
|
||||
fmax=7600, # To be increased/reduced depending on data.
|
||||
|
||||
###################### Our training parameters #################################
|
||||
img_size=96,
|
||||
fps=25,
|
||||
|
||||
batch_size=2,
|
||||
initial_learning_rate=1e-3,
|
||||
nepochs=100000, ### ctrl + c, stop whenever eval loss is consistently greater than train loss for ~10 epochs
|
||||
num_workers=0,
|
||||
checkpoint_interval=10000,
|
||||
eval_interval=10,
|
||||
writer_interval=5,
|
||||
save_optimizer_state=True,
|
||||
|
||||
syncnet_wt=0.0, # is initially zero, will be set automatically to 0.03 later. Leads to faster convergence.
|
||||
syncnet_batch_size=64,
|
||||
syncnet_lr=1e-4,
|
||||
syncnet_eval_interval=10000,
|
||||
syncnet_checkpoint_interval=10000,
|
||||
|
||||
disc_wt=0.07,
|
||||
disc_initial_learning_rate=1e-4,
|
||||
)
|
||||
|
||||
|
||||
def hparams_debug_string():
|
||||
values = hparams.values()
|
||||
hp = [" %s: %s" % (name, values[name]) for name in sorted(values) if name != "sentences"]
|
||||
return "Hyperparameters:\n" + "\n".join(hp)
|
||||
@@ -0,0 +1,47 @@
|
||||
import os
|
||||
import glob
|
||||
|
||||
def init_path(checkpoint_dir, config_dir, size=512, old_version=False, preprocess='crop'):
|
||||
|
||||
if old_version:
|
||||
#### load all the checkpoint of `pth`
|
||||
sadtalker_paths = {
|
||||
'wav2lip_checkpoint' : os.path.join(checkpoint_dir, 'wav2lip.pth'),
|
||||
'audio2pose_checkpoint' : os.path.join(checkpoint_dir, 'auido2pose_00140-model.pth'),
|
||||
'audio2exp_checkpoint' : os.path.join(checkpoint_dir, 'auido2exp_00300-model.pth'),
|
||||
'free_view_checkpoint' : os.path.join(checkpoint_dir, 'facevid2vid_00189-model.pth.tar'),
|
||||
'path_of_net_recon_model' : os.path.join(checkpoint_dir, 'epoch_20.pth')
|
||||
}
|
||||
|
||||
use_safetensor = False
|
||||
elif len(glob.glob(os.path.join(checkpoint_dir, '*.safetensors'))):
|
||||
print('using safetensor as default')
|
||||
sadtalker_paths = {
|
||||
"checkpoint":os.path.join(checkpoint_dir, 'SadTalker_V0.0.2_'+str(size)+'.safetensors'),
|
||||
}
|
||||
use_safetensor = True
|
||||
else:
|
||||
print("WARNING: The new version of the model will be updated by safetensor, you may need to download it mannully. We run the old version of the checkpoint this time!")
|
||||
use_safetensor = False
|
||||
|
||||
sadtalker_paths = {
|
||||
'wav2lip_checkpoint' : os.path.join(checkpoint_dir, 'wav2lip.pth'),
|
||||
'audio2pose_checkpoint' : os.path.join(checkpoint_dir, 'auido2pose_00140-model.pth'),
|
||||
'audio2exp_checkpoint' : os.path.join(checkpoint_dir, 'auido2exp_00300-model.pth'),
|
||||
'free_view_checkpoint' : os.path.join(checkpoint_dir, 'facevid2vid_00189-model.pth.tar'),
|
||||
'path_of_net_recon_model' : os.path.join(checkpoint_dir, 'epoch_20.pth')
|
||||
}
|
||||
|
||||
sadtalker_paths['dir_of_BFM_fitting'] = os.path.join(config_dir) # , 'BFM_Fitting'
|
||||
sadtalker_paths['audio2pose_yaml_path'] = os.path.join(config_dir, 'auido2pose.yaml')
|
||||
sadtalker_paths['audio2exp_yaml_path'] = os.path.join(config_dir, 'auido2exp.yaml')
|
||||
sadtalker_paths['use_safetensor'] = use_safetensor # os.path.join(config_dir, 'auido2exp.yaml')
|
||||
|
||||
if 'full' in preprocess:
|
||||
sadtalker_paths['mappingnet_checkpoint'] = os.path.join(checkpoint_dir, 'mapping_00109-model.pth.tar')
|
||||
sadtalker_paths['facerender_yaml'] = os.path.join(config_dir, 'facerender_still.yaml')
|
||||
else:
|
||||
sadtalker_paths['mappingnet_checkpoint'] = os.path.join(checkpoint_dir, 'mapping_00229-model.pth.tar')
|
||||
sadtalker_paths['facerender_yaml'] = os.path.join(config_dir, 'facerender.yaml')
|
||||
|
||||
return sadtalker_paths
|
||||
@@ -0,0 +1,141 @@
|
||||
import torch
|
||||
import yaml
|
||||
import os
|
||||
|
||||
import safetensors
|
||||
from safetensors.torch import save_file
|
||||
from yacs.config import CfgNode as CN
|
||||
import sys
|
||||
|
||||
sys.path.append('/apdcephfs/private_shadowcun/SadTalker')
|
||||
|
||||
from src.face3d.models import networks
|
||||
|
||||
from src.facerender.modules.keypoint_detector import HEEstimator, KPDetector
|
||||
from src.facerender.modules.mapping import MappingNet
|
||||
from src.facerender.modules.generator import OcclusionAwareGenerator, OcclusionAwareSPADEGenerator
|
||||
|
||||
from src.audio2pose_models.audio2pose import Audio2Pose
|
||||
from src.audio2exp_models.networks import SimpleWrapperV2
|
||||
from src.test_audio2coeff import load_cpk
|
||||
|
||||
size = 256
|
||||
############ face vid2vid
|
||||
config_path = os.path.join('src', 'config', 'facerender.yaml')
|
||||
current_root_path = '.'
|
||||
|
||||
path_of_net_recon_model = os.path.join(current_root_path, 'checkpoints', 'epoch_20.pth')
|
||||
net_recon = networks.define_net_recon(net_recon='resnet50', use_last_fc=False, init_path='')
|
||||
checkpoint = torch.load(path_of_net_recon_model, map_location='cpu')
|
||||
net_recon.load_state_dict(checkpoint['net_recon'])
|
||||
|
||||
with open(config_path) as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
generator = OcclusionAwareSPADEGenerator(**config['model_params']['generator_params'],
|
||||
**config['model_params']['common_params'])
|
||||
kp_extractor = KPDetector(**config['model_params']['kp_detector_params'],
|
||||
**config['model_params']['common_params'])
|
||||
he_estimator = HEEstimator(**config['model_params']['he_estimator_params'],
|
||||
**config['model_params']['common_params'])
|
||||
mapping = MappingNet(**config['model_params']['mapping_params'])
|
||||
|
||||
def load_cpk_facevid2vid(checkpoint_path, generator=None, discriminator=None,
|
||||
kp_detector=None, he_estimator=None, optimizer_generator=None,
|
||||
optimizer_discriminator=None, optimizer_kp_detector=None,
|
||||
optimizer_he_estimator=None, device="cpu"):
|
||||
|
||||
checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
|
||||
if generator is not None:
|
||||
generator.load_state_dict(checkpoint['generator'])
|
||||
if kp_detector is not None:
|
||||
kp_detector.load_state_dict(checkpoint['kp_detector'])
|
||||
if he_estimator is not None:
|
||||
he_estimator.load_state_dict(checkpoint['he_estimator'])
|
||||
if discriminator is not None:
|
||||
try:
|
||||
discriminator.load_state_dict(checkpoint['discriminator'])
|
||||
except:
|
||||
print ('No discriminator in the state-dict. Dicriminator will be randomly initialized')
|
||||
if optimizer_generator is not None:
|
||||
optimizer_generator.load_state_dict(checkpoint['optimizer_generator'])
|
||||
if optimizer_discriminator is not None:
|
||||
try:
|
||||
optimizer_discriminator.load_state_dict(checkpoint['optimizer_discriminator'])
|
||||
except RuntimeError as e:
|
||||
print ('No discriminator optimizer in the state-dict. Optimizer will be not initialized')
|
||||
if optimizer_kp_detector is not None:
|
||||
optimizer_kp_detector.load_state_dict(checkpoint['optimizer_kp_detector'])
|
||||
if optimizer_he_estimator is not None:
|
||||
optimizer_he_estimator.load_state_dict(checkpoint['optimizer_he_estimator'])
|
||||
|
||||
return checkpoint['epoch']
|
||||
|
||||
|
||||
def load_cpk_facevid2vid_safetensor(checkpoint_path, generator=None,
|
||||
kp_detector=None, he_estimator=None,
|
||||
device="cpu"):
|
||||
|
||||
checkpoint = safetensors.torch.load_file(checkpoint_path)
|
||||
|
||||
if generator is not None:
|
||||
x_generator = {}
|
||||
for k,v in checkpoint.items():
|
||||
if 'generator' in k:
|
||||
x_generator[k.replace('generator.', '')] = v
|
||||
generator.load_state_dict(x_generator)
|
||||
if kp_detector is not None:
|
||||
x_generator = {}
|
||||
for k,v in checkpoint.items():
|
||||
if 'kp_extractor' in k:
|
||||
x_generator[k.replace('kp_extractor.', '')] = v
|
||||
kp_detector.load_state_dict(x_generator)
|
||||
if he_estimator is not None:
|
||||
x_generator = {}
|
||||
for k,v in checkpoint.items():
|
||||
if 'he_estimator' in k:
|
||||
x_generator[k.replace('he_estimator.', '')] = v
|
||||
he_estimator.load_state_dict(x_generator)
|
||||
|
||||
return None
|
||||
|
||||
free_view_checkpoint = '/apdcephfs/private_shadowcun/SadTalker/checkpoints/facevid2vid_'+str(size)+'-model.pth.tar'
|
||||
load_cpk_facevid2vid(free_view_checkpoint, kp_detector=kp_extractor, generator=generator, he_estimator=he_estimator)
|
||||
|
||||
wav2lip_checkpoint = os.path.join(current_root_path, 'checkpoints', 'wav2lip.pth')
|
||||
|
||||
audio2pose_checkpoint = os.path.join(current_root_path, 'checkpoints', 'auido2pose_00140-model.pth')
|
||||
audio2pose_yaml_path = os.path.join(current_root_path, 'src', 'config', 'auido2pose.yaml')
|
||||
|
||||
audio2exp_checkpoint = os.path.join(current_root_path, 'checkpoints', 'auido2exp_00300-model.pth')
|
||||
audio2exp_yaml_path = os.path.join(current_root_path, 'src', 'config', 'auido2exp.yaml')
|
||||
|
||||
fcfg_pose = open(audio2pose_yaml_path)
|
||||
cfg_pose = CN.load_cfg(fcfg_pose)
|
||||
cfg_pose.freeze()
|
||||
audio2pose_model = Audio2Pose(cfg_pose, wav2lip_checkpoint)
|
||||
audio2pose_model.eval()
|
||||
load_cpk(audio2pose_checkpoint, model=audio2pose_model, device='cpu')
|
||||
|
||||
# load audio2exp_model
|
||||
netG = SimpleWrapperV2()
|
||||
netG.eval()
|
||||
load_cpk(audio2exp_checkpoint, model=netG, device='cpu')
|
||||
|
||||
class SadTalker(torch.nn.Module):
|
||||
def __init__(self, kp_extractor, generator, netG, audio2pose, face_3drecon):
|
||||
super(SadTalker, self).__init__()
|
||||
self.kp_extractor = kp_extractor
|
||||
self.generator = generator
|
||||
self.audio2exp = netG
|
||||
self.audio2pose = audio2pose
|
||||
self.face_3drecon = face_3drecon
|
||||
|
||||
|
||||
model = SadTalker(kp_extractor, generator, netG, audio2pose_model, net_recon)
|
||||
|
||||
# here, we want to convert it to safetensor
|
||||
save_file(model.state_dict(), "checkpoints/SadTalker_V0.0.2_"+str(size)+".safetensors")
|
||||
|
||||
### test
|
||||
load_cpk_facevid2vid_safetensor('checkpoints/SadTalker_V0.0.2_'+str(size)+'.safetensors', kp_detector=kp_extractor, generator=generator, he_estimator=None)
|
||||
@@ -0,0 +1,69 @@
|
||||
import cv2, os
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
import uuid
|
||||
|
||||
from src.utils.videoio import save_video_with_watermark
|
||||
|
||||
def paste_pic(video_path, pic_path, crop_info, new_audio_path, full_video_path, extended_crop=False):
|
||||
|
||||
if not os.path.isfile(pic_path):
|
||||
raise ValueError('pic_path must be a valid path to video/image file')
|
||||
elif pic_path.split('.')[-1] in ['jpg', 'png', 'jpeg']:
|
||||
# loader for first frame
|
||||
full_img = cv2.imread(pic_path)
|
||||
else:
|
||||
# loader for videos
|
||||
video_stream = cv2.VideoCapture(pic_path)
|
||||
fps = video_stream.get(cv2.CAP_PROP_FPS)
|
||||
full_frames = []
|
||||
while 1:
|
||||
still_reading, frame = video_stream.read()
|
||||
if not still_reading:
|
||||
video_stream.release()
|
||||
break
|
||||
break
|
||||
full_img = frame
|
||||
frame_h = full_img.shape[0]
|
||||
frame_w = full_img.shape[1]
|
||||
|
||||
video_stream = cv2.VideoCapture(video_path)
|
||||
fps = video_stream.get(cv2.CAP_PROP_FPS)
|
||||
crop_frames = []
|
||||
while 1:
|
||||
still_reading, frame = video_stream.read()
|
||||
if not still_reading:
|
||||
video_stream.release()
|
||||
break
|
||||
crop_frames.append(frame)
|
||||
|
||||
if len(crop_info) != 3:
|
||||
print("you didn't crop the image")
|
||||
return
|
||||
else:
|
||||
r_w, r_h = crop_info[0]
|
||||
clx, cly, crx, cry = crop_info[1]
|
||||
lx, ly, rx, ry = crop_info[2]
|
||||
lx, ly, rx, ry = int(lx), int(ly), int(rx), int(ry)
|
||||
# oy1, oy2, ox1, ox2 = cly+ly, cly+ry, clx+lx, clx+rx
|
||||
# oy1, oy2, ox1, ox2 = cly+ly, cly+ry, clx+lx, clx+rx
|
||||
|
||||
if extended_crop:
|
||||
oy1, oy2, ox1, ox2 = cly, cry, clx, crx
|
||||
else:
|
||||
oy1, oy2, ox1, ox2 = cly+ly, cly+ry, clx+lx, clx+rx
|
||||
|
||||
tmp_path = str(uuid.uuid4())+'.mp4'
|
||||
out_tmp = cv2.VideoWriter(tmp_path, cv2.VideoWriter_fourcc(*'MP4V'), fps, (frame_w, frame_h))
|
||||
for crop_frame in tqdm(crop_frames, 'seamlessClone:'):
|
||||
p = cv2.resize(crop_frame.astype(np.uint8), (ox2-ox1, oy2 - oy1))
|
||||
|
||||
mask = 255*np.ones(p.shape, p.dtype)
|
||||
location = ((ox1+ox2) // 2, (oy1+oy2) // 2)
|
||||
gen_img = cv2.seamlessClone(p, full_img, mask, location, cv2.NORMAL_CLONE)
|
||||
out_tmp.write(gen_img)
|
||||
|
||||
out_tmp.release()
|
||||
|
||||
save_video_with_watermark(tmp_path, new_audio_path, full_video_path, watermark=False)
|
||||
os.remove(tmp_path)
|
||||
@@ -0,0 +1,170 @@
|
||||
import numpy as np
|
||||
import cv2, os, sys, torch
|
||||
from tqdm import tqdm
|
||||
from PIL import Image
|
||||
|
||||
# 3dmm extraction
|
||||
import safetensors
|
||||
import safetensors.torch
|
||||
from src.face3d.util.preprocess import align_img
|
||||
from src.face3d.util.load_mats import load_lm3d
|
||||
from src.face3d.models import networks
|
||||
|
||||
from scipy.io import loadmat, savemat
|
||||
from src.utils.croper import Preprocesser
|
||||
|
||||
|
||||
import warnings
|
||||
|
||||
from src.utils.safetensor_helper import load_x_from_safetensor
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
def split_coeff(coeffs):
|
||||
"""
|
||||
Return:
|
||||
coeffs_dict -- a dict of torch.tensors
|
||||
|
||||
Parameters:
|
||||
coeffs -- torch.tensor, size (B, 256)
|
||||
"""
|
||||
id_coeffs = coeffs[:, :80]
|
||||
exp_coeffs = coeffs[:, 80: 144]
|
||||
tex_coeffs = coeffs[:, 144: 224]
|
||||
angles = coeffs[:, 224: 227]
|
||||
gammas = coeffs[:, 227: 254]
|
||||
translations = coeffs[:, 254:]
|
||||
return {
|
||||
'id': id_coeffs,
|
||||
'exp': exp_coeffs,
|
||||
'tex': tex_coeffs,
|
||||
'angle': angles,
|
||||
'gamma': gammas,
|
||||
'trans': translations
|
||||
}
|
||||
|
||||
|
||||
class CropAndExtract():
|
||||
def __init__(self, sadtalker_path, device):
|
||||
|
||||
self.propress = Preprocesser(device)
|
||||
self.net_recon = networks.define_net_recon(net_recon='resnet50', use_last_fc=False, init_path='').to(device)
|
||||
|
||||
if sadtalker_path['use_safetensor']:
|
||||
checkpoint = safetensors.torch.load_file(sadtalker_path['checkpoint'])
|
||||
self.net_recon.load_state_dict(load_x_from_safetensor(checkpoint, 'face_3drecon'))
|
||||
else:
|
||||
checkpoint = torch.load(sadtalker_path['path_of_net_recon_model'], map_location=torch.device(device))
|
||||
self.net_recon.load_state_dict(checkpoint['net_recon'])
|
||||
|
||||
self.net_recon.eval()
|
||||
self.lm3d_std = load_lm3d(sadtalker_path['dir_of_BFM_fitting'])
|
||||
self.device = device
|
||||
|
||||
def generate(self, input_path, save_dir, crop_or_resize='crop', source_image_flag=False, pic_size=256):
|
||||
|
||||
pic_name = os.path.splitext(os.path.split(input_path)[-1])[0]
|
||||
|
||||
landmarks_path = os.path.join(save_dir, pic_name+'_landmarks.txt')
|
||||
coeff_path = os.path.join(save_dir, pic_name+'.mat')
|
||||
png_path = os.path.join(save_dir, pic_name+'.png')
|
||||
|
||||
#load input
|
||||
if not os.path.isfile(input_path):
|
||||
raise ValueError('input_path must be a valid path to video/image file')
|
||||
elif input_path.split('.')[-1] in ['jpg', 'png', 'jpeg']:
|
||||
# loader for first frame
|
||||
full_frames = [cv2.imread(input_path)]
|
||||
fps = 25
|
||||
else:
|
||||
# loader for videos
|
||||
video_stream = cv2.VideoCapture(input_path)
|
||||
fps = video_stream.get(cv2.CAP_PROP_FPS)
|
||||
full_frames = []
|
||||
while 1:
|
||||
still_reading, frame = video_stream.read()
|
||||
if not still_reading:
|
||||
video_stream.release()
|
||||
break
|
||||
full_frames.append(frame)
|
||||
if source_image_flag:
|
||||
break
|
||||
|
||||
x_full_frames= [cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in full_frames]
|
||||
|
||||
#### crop images as the
|
||||
if 'crop' in crop_or_resize.lower(): # default crop
|
||||
x_full_frames, crop, quad = self.propress.crop(x_full_frames, still=True if 'ext' in crop_or_resize.lower() else False, xsize=512)
|
||||
clx, cly, crx, cry = crop
|
||||
lx, ly, rx, ry = quad
|
||||
lx, ly, rx, ry = int(lx), int(ly), int(rx), int(ry)
|
||||
oy1, oy2, ox1, ox2 = cly+ly, cly+ry, clx+lx, clx+rx
|
||||
crop_info = ((ox2 - ox1, oy2 - oy1), crop, quad)
|
||||
elif 'full' in crop_or_resize.lower():
|
||||
x_full_frames, crop, quad = self.propress.crop(x_full_frames, still=True if 'ext' in crop_or_resize.lower() else False, xsize=512)
|
||||
clx, cly, crx, cry = crop
|
||||
lx, ly, rx, ry = quad
|
||||
lx, ly, rx, ry = int(lx), int(ly), int(rx), int(ry)
|
||||
oy1, oy2, ox1, ox2 = cly+ly, cly+ry, clx+lx, clx+rx
|
||||
crop_info = ((ox2 - ox1, oy2 - oy1), crop, quad)
|
||||
else: # resize mode
|
||||
oy1, oy2, ox1, ox2 = 0, x_full_frames[0].shape[0], 0, x_full_frames[0].shape[1]
|
||||
crop_info = ((ox2 - ox1, oy2 - oy1), None, None)
|
||||
|
||||
frames_pil = [Image.fromarray(cv2.resize(frame,(pic_size, pic_size))) for frame in x_full_frames]
|
||||
if len(frames_pil) == 0:
|
||||
print('No face is detected in the input file')
|
||||
return None, None
|
||||
|
||||
# save crop info
|
||||
for frame in frames_pil:
|
||||
cv2.imwrite(png_path, cv2.cvtColor(np.array(frame), cv2.COLOR_RGB2BGR))
|
||||
|
||||
# 2. get the landmark according to the detected face.
|
||||
if not os.path.isfile(landmarks_path):
|
||||
lm = self.propress.predictor.extract_keypoint(frames_pil, landmarks_path)
|
||||
else:
|
||||
print(' Using saved landmarks.')
|
||||
lm = np.loadtxt(landmarks_path).astype(np.float32)
|
||||
lm = lm.reshape([len(x_full_frames), -1, 2])
|
||||
|
||||
if not os.path.isfile(coeff_path):
|
||||
# load 3dmm paramter generator from Deep3DFaceRecon_pytorch
|
||||
video_coeffs, full_coeffs = [], []
|
||||
for idx in tqdm(range(len(frames_pil)), desc='3DMM Extraction In Video:'):
|
||||
frame = frames_pil[idx]
|
||||
W,H = frame.size
|
||||
lm1 = lm[idx].reshape([-1, 2])
|
||||
|
||||
if np.mean(lm1) == -1:
|
||||
lm1 = (self.lm3d_std[:, :2]+1)/2.
|
||||
lm1 = np.concatenate(
|
||||
[lm1[:, :1]*W, lm1[:, 1:2]*H], 1
|
||||
)
|
||||
else:
|
||||
lm1[:, -1] = H - 1 - lm1[:, -1]
|
||||
|
||||
trans_params, im1, lm1, _ = align_img(frame, lm1, self.lm3d_std)
|
||||
|
||||
trans_params = np.array([float(item) for item in np.hsplit(trans_params, 5)]).astype(np.float32)
|
||||
im_t = torch.tensor(np.array(im1)/255., dtype=torch.float32).permute(2, 0, 1).to(self.device).unsqueeze(0)
|
||||
|
||||
with torch.no_grad():
|
||||
full_coeff = self.net_recon(im_t)
|
||||
coeffs = split_coeff(full_coeff)
|
||||
|
||||
pred_coeff = {key:coeffs[key].cpu().numpy() for key in coeffs}
|
||||
|
||||
pred_coeff = np.concatenate([
|
||||
pred_coeff['exp'],
|
||||
pred_coeff['angle'],
|
||||
pred_coeff['trans'],
|
||||
trans_params[2:][None],
|
||||
], 1)
|
||||
video_coeffs.append(pred_coeff)
|
||||
full_coeffs.append(full_coeff.cpu().numpy())
|
||||
|
||||
semantic_npy = np.array(video_coeffs)[:,0]
|
||||
|
||||
savemat(coeff_path, {'coeff_3dmm': semantic_npy, 'full_3dmm': np.array(full_coeffs)[0]})
|
||||
|
||||
return coeff_path, png_path, crop_info
|
||||
@@ -0,0 +1,8 @@
|
||||
|
||||
|
||||
def load_x_from_safetensor(checkpoint, key):
|
||||
x_generator = {}
|
||||
for k,v in checkpoint.items():
|
||||
if key in k:
|
||||
x_generator[k.replace(key+'.', '')] = v
|
||||
return x_generator
|
||||
@@ -0,0 +1,20 @@
|
||||
import os
|
||||
import tempfile
|
||||
from TTS.api import TTS
|
||||
|
||||
|
||||
class TTSTalker():
|
||||
def __init__(self) -> None:
|
||||
model_name = TTS().list_models()[0]
|
||||
self.tts = TTS(model_name)
|
||||
|
||||
def test(self, text, language='en'):
|
||||
|
||||
tempf = tempfile.NamedTemporaryFile(
|
||||
delete = False,
|
||||
suffix = ('.'+'wav'),
|
||||
)
|
||||
|
||||
self.tts.tts_to_file(text, speaker=self.tts.speakers[0], language=language, file_path=tempf.name)
|
||||
|
||||
return tempf.name
|
||||
@@ -0,0 +1,41 @@
|
||||
import shutil
|
||||
import uuid
|
||||
|
||||
import os
|
||||
|
||||
import cv2
|
||||
|
||||
def load_video_to_cv2(input_path):
|
||||
video_stream = cv2.VideoCapture(input_path)
|
||||
fps = video_stream.get(cv2.CAP_PROP_FPS)
|
||||
full_frames = []
|
||||
while 1:
|
||||
still_reading, frame = video_stream.read()
|
||||
if not still_reading:
|
||||
video_stream.release()
|
||||
break
|
||||
full_frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
return full_frames
|
||||
|
||||
def save_video_with_watermark(video, audio, save_path, watermark=False):
|
||||
temp_file = str(uuid.uuid4())+'.mp4'
|
||||
cmd = r'ffmpeg -y -hide_banner -loglevel error -i "%s" -i "%s" -vcodec copy "%s"' % (video, audio, temp_file)
|
||||
os.system(cmd)
|
||||
|
||||
if watermark is False:
|
||||
shutil.move(temp_file, save_path)
|
||||
else:
|
||||
# watermark
|
||||
try:
|
||||
##### check if stable-diffusion-webui
|
||||
import webui
|
||||
from modules import paths
|
||||
watarmark_path = paths.script_path+"/extensions/SadTalker/docs/sadtalker_logo.png"
|
||||
except:
|
||||
# get the root path of sadtalker.
|
||||
dir_path = os.path.dirname(os.path.realpath(__file__))
|
||||
watarmark_path = dir_path+"/../../docs/sadtalker_logo.png"
|
||||
|
||||
cmd = r'ffmpeg -y -hide_banner -loglevel error -i "%s" -i "%s" -filter_complex "[1]scale=100:-1[wm];[0][wm]overlay=(main_w-overlay_w)-10:10" "%s"' % (temp_file, watarmark_path, save_path)
|
||||
os.system(cmd)
|
||||
os.remove(temp_file)
|
||||
Reference in New Issue
Block a user