import numpy as np import cv2, os, sys, subprocess, platform, torch from tqdm import tqdm from PIL import Image from scipy.io import loadmat sys.path.insert(0, 'third_part') sys.path.insert(0, 'third_part/GPEN') sys.path.insert(0, 'third_part/GFPGAN') # 3dmm extraction from third_part.face3d.util.preprocess import align_img from third_part.face3d.util.load_mats import load_lm3d from third_part.face3d.extract_kp_videos import KeypointExtractor # face enhancement from third_part.GPEN.gpen_face_enhancer import FaceEnhancement from third_part.GFPGAN.gfpgan import GFPGANer # expression control from third_part.ganimation_replicate.model.ganimation import GANimationModel from utils import audio from utils.ffhq_preprocess import Croper from utils.alignment_stit import crop_faces, calc_alignment_coefficients, paste_image from utils.inference_utils import Laplacian_Pyramid_Blending_with_mask, face_detect, load_model, options, split_coeff, \ trans_image, transform_semantic, find_crop_norm_ratio, load_face3d_net, exp_aus_dict import warnings warnings.filterwarnings("ignore") args = options() def main(): device = 'cuda' if torch.cuda.is_available() else 'cpu' print('[Info] Using {} for inference.'.format(device)) os.makedirs(os.path.join('temp', args.tmp_dir), exist_ok=True) enhancer = FaceEnhancement(base_dir='checkpoints', size=512, model='GPEN-BFR-512', use_sr=False, \ sr_model='rrdb_realesrnet_psnr', channel_multiplier=2, narrow=1, device=device) restorer = GFPGANer(model_path='checkpoints/GFPGANv1.3.pth', upscale=1, arch='clean', \ channel_multiplier=2, bg_upsampler=None) base_name = args.face.split('/')[-1] if os.path.isfile(args.face) and args.face.split('.')[1] in ['jpg', 'png', 'jpeg']: args.static = True if not os.path.isfile(args.face): raise ValueError('--face argument must be a valid path to video/image file') elif args.face.split('.')[1] in ['jpg', 'png', 'jpeg']: full_frames = [cv2.imread(args.face)] fps = args.fps else: video_stream = cv2.VideoCapture(args.face) fps = video_stream.get(cv2.CAP_PROP_FPS) full_frames = [] while True: still_reading, frame = video_stream.read() if not still_reading: video_stream.release() break y1, y2, x1, x2 = args.crop if x2 == -1: x2 = frame.shape[1] if y2 == -1: y2 = frame.shape[0] frame = frame[y1:y2, x1:x2] full_frames.append(frame) print ("[Step 0] Number of frames available for inference: "+str(len(full_frames))) # face detection & cropping, cropping the first frame as the style of FFHQ croper = Croper('checkpoints/shape_predictor_68_face_landmarks.dat') full_frames_RGB = [cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in full_frames] full_frames_RGB, crop, quad = croper.crop(full_frames_RGB, 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, min(cly+ry, full_frames[0].shape[0]), clx+lx, min(clx+rx, full_frames[0].shape[1]) # original_size = (ox2 - ox1, oy2 - oy1) frames_pil = [Image.fromarray(cv2.resize(frame,(256,256))) for frame in full_frames_RGB] # get the landmark according to the detected face. if not os.path.isfile('temp/'+base_name+'_landmarks.txt') or args.re_preprocess: print('[Step 1] Landmarks Extraction in Video.') kp_extractor = KeypointExtractor() lm = kp_extractor.extract_keypoint(frames_pil, './temp/'+base_name+'_landmarks.txt') else: print('[Step 1] Using saved landmarks.') lm = np.loadtxt('temp/'+base_name+'_landmarks.txt').astype(np.float32) lm = lm.reshape([len(full_frames), -1, 2]) if not os.path.isfile('temp/'+base_name+'_coeffs.npy') or args.exp_img is not None or args.re_preprocess: net_recon = load_face3d_net(args.face3d_net_path, device) lm3d_std = load_lm3d('checkpoints/BFM') video_coeffs = [] for idx in tqdm(range(len(frames_pil)), desc="[Step 2] 3DMM Extraction In Video:"): frame = frames_pil[idx] W, H = frame.size lm_idx = lm[idx].reshape([-1, 2]) if np.mean(lm_idx) == -1: lm_idx = (lm3d_std[:, :2]+1) / 2. lm_idx = np.concatenate([lm_idx[:, :1] * W, lm_idx[:, 1:2] * H], 1) else: lm_idx[:, -1] = H - 1 - lm_idx[:, -1] trans_params, im_idx, lm_idx, _ = align_img(frame, lm_idx, lm3d_std) trans_params = np.array([float(item) for item in np.hsplit(trans_params, 5)]).astype(np.float32) im_idx_tensor = torch.tensor(np.array(im_idx)/255., dtype=torch.float32).permute(2, 0, 1).to(device).unsqueeze(0) with torch.no_grad(): coeffs = split_coeff(net_recon(im_idx_tensor)) pred_coeff = {key:coeffs[key].cpu().numpy() for key in coeffs} pred_coeff = np.concatenate([pred_coeff['id'], pred_coeff['exp'], pred_coeff['tex'], pred_coeff['angle'],\ pred_coeff['gamma'], pred_coeff['trans'], trans_params[None]], 1) video_coeffs.append(pred_coeff) semantic_npy = np.array(video_coeffs)[:,0] np.save('temp/'+base_name+'_coeffs.npy', semantic_npy) else: print('[Step 2] Using saved coeffs.') semantic_npy = np.load('temp/'+base_name+'_coeffs.npy').astype(np.float32) # generate the 3dmm coeff from a single image if args.exp_img is not None and ('.png' in args.exp_img or '.jpg' in args.exp_img): print('extract the exp from',args.exp_img) exp_pil = Image.open(args.exp_img).convert('RGB') lm3d_std = load_lm3d('third_part/face3d/BFM') W, H = exp_pil.size kp_extractor = KeypointExtractor() lm_exp = kp_extractor.extract_keypoint([exp_pil], 'temp/'+base_name+'_temp.txt')[0] if np.mean(lm_exp) == -1: lm_exp = (lm3d_std[:, :2] + 1) / 2. lm_exp = np.concatenate( [lm_exp[:, :1] * W, lm_exp[:, 1:2] * H], 1) else: lm_exp[:, -1] = H - 1 - lm_exp[:, -1] trans_params, im_exp, lm_exp, _ = align_img(exp_pil, lm_exp, lm3d_std) trans_params = np.array([float(item) for item in np.hsplit(trans_params, 5)]).astype(np.float32) im_exp_tensor = torch.tensor(np.array(im_exp)/255., dtype=torch.float32).permute(2, 0, 1).to(device).unsqueeze(0) with torch.no_grad(): expression = split_coeff(net_recon(im_exp_tensor))['exp'][0] del net_recon elif args.exp_img == 'smile': expression = torch.tensor(loadmat('checkpoints/expression.mat')['expression_mouth'])[0] else: print('using expression center') expression = torch.tensor(loadmat('checkpoints/expression.mat')['expression_center'])[0] # load DNet, model(LNet and ENet) D_Net, model = load_model(args, device) if not os.path.isfile('temp/'+base_name+'_stablized.npy') or args.re_preprocess: imgs = [] for idx in tqdm(range(len(frames_pil)), desc="[Step 3] Stablize the expression In Video:"): if args.one_shot: source_img = trans_image(frames_pil[0]).unsqueeze(0).to(device) semantic_source_numpy = semantic_npy[0:1] else: source_img = trans_image(frames_pil[idx]).unsqueeze(0).to(device) semantic_source_numpy = semantic_npy[idx:idx+1] ratio = find_crop_norm_ratio(semantic_source_numpy, semantic_npy) coeff = transform_semantic(semantic_npy, idx, ratio).unsqueeze(0).to(device) # hacking the new expression coeff[:, :64, :] = expression[None, :64, None].to(device) with torch.no_grad(): output = D_Net(source_img, coeff) img_stablized = np.uint8((output['fake_image'].squeeze(0).permute(1,2,0).cpu().clamp_(-1, 1).numpy() + 1 )/2. * 255) imgs.append(cv2.cvtColor(img_stablized,cv2.COLOR_RGB2BGR)) np.save('temp/'+base_name+'_stablized.npy',imgs) del D_Net else: print('[Step 3] Using saved stablized video.') imgs = np.load('temp/'+base_name+'_stablized.npy') torch.cuda.empty_cache() if not args.audio.endswith('.wav'): command = 'ffmpeg -loglevel error -y -i {} -strict -2 {}'.format(args.audio, 'temp/{}/temp.wav'.format(args.tmp_dir)) subprocess.call(command, shell=True) args.audio = 'temp/{}/temp.wav'.format(args.tmp_dir) wav = audio.load_wav(args.audio, 16000) mel = audio.melspectrogram(wav) if np.isnan(mel.reshape(-1)).sum() > 0: raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again') mel_step_size, mel_idx_multiplier, i, mel_chunks = 16, 80./fps, 0, [] while True: start_idx = int(i * mel_idx_multiplier) if start_idx + mel_step_size > len(mel[0]): mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:]) break mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size]) i += 1 print("[Step 4] Load audio; Length of mel chunks: {}".format(len(mel_chunks))) imgs = imgs[:len(mel_chunks)] full_frames = full_frames[:len(mel_chunks)] lm = lm[:len(mel_chunks)] imgs_enhanced = [] for idx in tqdm(range(len(imgs)), desc='[Step 5] Reference Enhancement'): img = imgs[idx] pred, _, _ = enhancer.process(img, img, face_enhance=True, possion_blending=False) imgs_enhanced.append(pred) gen = datagen(imgs_enhanced.copy(), mel_chunks, full_frames, None, (oy1,oy2,ox1,ox2)) frame_h, frame_w = full_frames[0].shape[:-1] out = cv2.VideoWriter('temp/{}/result.mp4'.format(args.tmp_dir), cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_w, frame_h)) if args.up_face != 'original': instance = GANimationModel() instance.initialize() instance.setup() kp_extractor = KeypointExtractor() for i, (img_batch, mel_batch, frames, coords, img_original, f_frames) in enumerate(tqdm(gen, desc='[Step 6] Lip Synthesis:', total=int(np.ceil(float(len(mel_chunks)) / args.LNet_batch_size)))): img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device) mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device) img_original = torch.FloatTensor(np.transpose(img_original, (0, 3, 1, 2))).to(device)/255. # BGR -> RGB with torch.no_grad(): incomplete, reference = torch.split(img_batch, 3, dim=1) pred, low_res = model(mel_batch, img_batch, reference) pred = torch.clamp(pred, 0, 1) if args.up_face in ['sad', 'angry', 'surprise']: tar_aus = exp_aus_dict[args.up_face] else: pass if args.up_face == 'original': cur_gen_faces = img_original else: test_batch = {'src_img': torch.nn.functional.interpolate((img_original * 2 - 1), size=(128, 128), mode='bilinear'), 'tar_aus': tar_aus.repeat(len(incomplete), 1)} instance.feed_batch(test_batch) instance.forward() cur_gen_faces = torch.nn.functional.interpolate(instance.fake_img / 2. + 0.5, size=(384, 384), mode='bilinear') if args.without_rl1 is not False: incomplete, reference = torch.split(img_batch, 3, dim=1) mask = torch.where(incomplete==0, torch.ones_like(incomplete), torch.zeros_like(incomplete)) pred = pred * mask + cur_gen_faces * (1 - mask) pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255. torch.cuda.empty_cache() for p, f, xf, c in zip(pred, frames, f_frames, coords): y1, y2, x1, x2 = c p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1)) ff = xf.copy() ff[y1:y2, x1:x2] = p # month region enhancement by GFPGAN cropped_faces, restored_faces, restored_img = restorer.enhance( ff, has_aligned=False, only_center_face=True, paste_back=True) # 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, mm = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 255, 255, 255, 0, 0, 0, 0, 0, 0] mouse_mask = np.zeros_like(restored_img) tmp_mask = enhancer.faceparser.process(restored_img[y1:y2, x1:x2], mm)[0] mouse_mask[y1:y2, x1:x2]= cv2.resize(tmp_mask, (x2 - x1, y2 - y1))[:, :, np.newaxis] / 255. height, width = ff.shape[:2] restored_img, ff, full_mask = [cv2.resize(x, (512, 512)) for x in (restored_img, ff, np.float32(mouse_mask))] img = Laplacian_Pyramid_Blending_with_mask(restored_img, ff, full_mask[:, :, 0], 10) pp = np.uint8(cv2.resize(np.clip(img, 0 ,255), (width, height))) pp, orig_faces, enhanced_faces = enhancer.process(pp, xf, bbox=c, face_enhance=False, possion_blending=True) out.write(pp) out.release() if not os.path.isdir(os.path.dirname(args.outfile)): os.makedirs(os.path.dirname(args.outfile), exist_ok=True) command = 'ffmpeg -loglevel error -y -i {} -i {} -strict -2 -q:v 1 {}'.format(args.audio, 'temp/{}/result.mp4'.format(args.tmp_dir), args.outfile) subprocess.call(command, shell=platform.system() != 'Windows') print('outfile:', args.outfile) # frames:256x256, full_frames: original size def datagen(frames, mels, full_frames, frames_pil, cox): img_batch, mel_batch, frame_batch, coords_batch, ref_batch, full_frame_batch = [], [], [], [], [], [] base_name = args.face.split('/')[-1] refs = [] image_size = 256 # original frames kp_extractor = KeypointExtractor() fr_pil = [Image.fromarray(frame) for frame in frames] lms = kp_extractor.extract_keypoint(fr_pil, 'temp/'+base_name+'x12_landmarks.txt') frames_pil = [ (lm, frame) for frame,lm in zip(fr_pil, lms)] # frames is the croped version of modified face crops, orig_images, quads = crop_faces(image_size, frames_pil, scale=1.0, use_fa=True) inverse_transforms = [calc_alignment_coefficients(quad + 0.5, [[0, 0], [0, image_size], [image_size, image_size], [image_size, 0]]) for quad in quads] del kp_extractor.detector oy1,oy2,ox1,ox2 = cox face_det_results = face_detect(full_frames, args, jaw_correction=True) for inverse_transform, crop, full_frame, face_det in zip(inverse_transforms, crops, full_frames, face_det_results): imc_pil = paste_image(inverse_transform, crop, Image.fromarray( cv2.resize(full_frame[int(oy1):int(oy2), int(ox1):int(ox2)], (256, 256)))) ff = full_frame.copy() ff[int(oy1):int(oy2), int(ox1):int(ox2)] = cv2.resize(np.array(imc_pil.convert('RGB')), (ox2 - ox1, oy2 - oy1)) oface, coords = face_det y1, y2, x1, x2 = coords refs.append(ff[y1: y2, x1:x2]) for i, m in enumerate(mels): idx = 0 if args.static else i % len(frames) frame_to_save = frames[idx].copy() face = refs[idx] oface, coords = face_det_results[idx].copy() face = cv2.resize(face, (args.img_size, args.img_size)) oface = cv2.resize(oface, (args.img_size, args.img_size)) img_batch.append(oface) ref_batch.append(face) mel_batch.append(m) coords_batch.append(coords) frame_batch.append(frame_to_save) full_frame_batch.append(full_frames[idx].copy()) if len(img_batch) >= args.LNet_batch_size: img_batch, mel_batch, ref_batch = np.asarray(img_batch), np.asarray(mel_batch), np.asarray(ref_batch) img_masked = img_batch.copy() img_original = img_batch.copy() img_masked[:, args.img_size//2:] = 0 img_batch = np.concatenate((img_masked, ref_batch), axis=3) / 255. mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) yield img_batch, mel_batch, frame_batch, coords_batch, img_original, full_frame_batch img_batch, mel_batch, frame_batch, coords_batch, img_original, full_frame_batch, ref_batch = [], [], [], [], [], [], [] if len(img_batch) > 0: img_batch, mel_batch, ref_batch = np.asarray(img_batch), np.asarray(mel_batch), np.asarray(ref_batch) img_masked = img_batch.copy() img_original = img_batch.copy() img_masked[:, args.img_size//2:] = 0 img_batch = np.concatenate((img_masked, ref_batch), axis=3) / 255. mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) yield img_batch, mel_batch, frame_batch, coords_batch, img_original, full_frame_batch if __name__ == '__main__': main()