|
|
|
import shutil
|
|
import cv2
|
|
import numpy as np
|
|
from src.face3d.models.bfm import ParametricFaceModel
|
|
from src.face3d.models.facerecon_model import FaceReconModel
|
|
import torch
|
|
import subprocess, platform
|
|
import scipy.io as scio
|
|
from tqdm import tqdm
|
|
|
|
|
|
def draw_landmarks(image, landmarks):
|
|
for i, point in enumerate(landmarks):
|
|
cv2.circle(image, (int(point[0]), int(point[1])), 2, (0, 255, 0), -1)
|
|
cv2.putText(image, str(i), (int(point[0]), int(point[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 255, 255), 1)
|
|
return image
|
|
|
|
|
|
def gen_composed_video(args, device, first_frame_coeff, coeff_path, audio_path, save_path, save_lmk_path, crop_info, extended_crop = False):
|
|
|
|
coeff_first = scio.loadmat(first_frame_coeff)['full_3dmm']
|
|
info = scio.loadmat(first_frame_coeff)['trans_params'][0]
|
|
print(info)
|
|
|
|
coeff_pred = scio.loadmat(coeff_path)['coeff_3dmm']
|
|
|
|
coeff_full = np.repeat(coeff_first, coeff_pred.shape[0], axis=0)
|
|
|
|
coeff_full[:, 80:144] = coeff_pred[:, 0:64]
|
|
coeff_full[:, 224:227] = coeff_pred[:, 64:67]
|
|
coeff_full[:, 254:] = coeff_pred[:, 67:]
|
|
|
|
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)
|
|
|
|
|
|
|
|
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_video_path = '/tmp/face3dtmp.mp4'
|
|
facemodel = FaceReconModel(args)
|
|
im0 = cv2.imread(args.source_image)
|
|
|
|
video = cv2.VideoWriter(tmp_video_path, cv2.VideoWriter_fourcc(*'mp4v'), 25, (224, 224))
|
|
|
|
|
|
|
|
x_scale, y_scale = (ox2 - ox1)/256 , (oy2 - oy1)/256
|
|
|
|
W, H = im0.shape[0], im0.shape[1]
|
|
|
|
_, _, s, _, _, orig_left, orig_up, orig_crop_size =(info[0], info[1], info[2], info[3], info[4], info[5], info[6], info[7])
|
|
orig_left, orig_up, orig_crop_size = [int(x) for x in (orig_left, orig_up, orig_crop_size)]
|
|
|
|
landmark_scale = np.array([[x_scale, y_scale]])
|
|
landmark_shift = np.array([[orig_left, orig_up]])
|
|
landmark_shift2 = np.array([[ox1, oy1]])
|
|
|
|
landmarks = []
|
|
|
|
print(orig_up, orig_left, orig_crop_size, s)
|
|
|
|
for k in tqdm(range(coeff_pred.shape[0]), 'face3d rendering:'):
|
|
cur_coeff_full = torch.tensor(coeff_full[k:k+1], device=device)
|
|
|
|
facemodel.forward(cur_coeff_full, device)
|
|
|
|
predicted_landmark = facemodel.pred_lm
|
|
predicted_landmark = predicted_landmark.cpu().numpy().squeeze()
|
|
|
|
predicted_landmark[:, 1] = 224 - predicted_landmark[:, 1]
|
|
|
|
predicted_landmark = ((predicted_landmark + landmark_shift) / s[0] * landmark_scale) + landmark_shift2
|
|
|
|
landmarks.append(predicted_landmark)
|
|
|
|
rendered_img = facemodel.pred_face
|
|
rendered_img = 255. * rendered_img.cpu().numpy().squeeze().transpose(1,2,0)
|
|
out_img = rendered_img[:, :, :3].astype(np.uint8)
|
|
|
|
video.write(np.uint8(out_img[:,:,::-1]))
|
|
|
|
video.release()
|
|
|
|
|
|
video = cv2.VideoWriter(save_lmk_path, cv2.VideoWriter_fourcc(*'mp4v'), 25, (im0.shape[0], im0.shape[1]))
|
|
|
|
for k in tqdm(range(len(landmarks)), 'face3d vis:'):
|
|
|
|
im = draw_landmarks(np.uint8(np.ones_like(im0)*255), landmarks[k])
|
|
video.write(im)
|
|
video.release()
|
|
|
|
shutil.copyfile(args.source_image, save_lmk_path.replace('.mp4', '.png'))
|
|
|
|
np.save(save_lmk_path.replace('.mp4', '.npy'), landmarks)
|
|
|
|
command = 'ffmpeg -v quiet -y -i {} -i {} -strict -2 -q:v 1 {}'.format(audio_path, tmp_video_path, save_path)
|
|
subprocess.call(command, shell=platform.system() != 'Windows')
|
|
|
|
|
