File size: 7,277 Bytes
d3fb87a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 |
import numpy as np
import cv2, os, sys, torch
from tqdm import tqdm
from PIL import Image
import time
# 3dmm extraction
import safetensors
import safetensors.torch
from chat_anything.sad_talker.face3d.util.preprocess import align_img
from chat_anything.sad_talker.face3d.util.load_mats import load_lm3d
from chat_anything.sad_talker.face3d.models import networks
from scipy.io import loadmat, savemat
from chat_anything.sad_talker.utils.croper import Preprocesser
import warnings
from chat_anything.sad_talker.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])
print(len(frames_pil))
print(frames_pil[0].size)
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)
start_time=time.time()
with torch.no_grad():
full_coeff = self.net_recon(im_t)
coeffs = split_coeff(full_coeff)
print(type(coeffs))
end_time=time.time()
ext_time=end_time-start_time
print("3DMM检测时间：%.4f秒" % ext_time)
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
|