606 lines
33 KiB
Python
606 lines
33 KiB
Python
import os,cv2,torch,subprocess,platform
|
|
import mediapipe as mp
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
from .draw_landmark import draw_landmarks
|
|
import face_alignment
|
|
from .face_mask import FaceMask
|
|
from cuda_malloc import cuda_malloc_supported
|
|
from .models import Landmark_generator as Landmark_transformer,Renderer,audio
|
|
|
|
NAME = "IP_LAP"
|
|
|
|
# the following is the index sequence for fical landmarks detected by mediapipe
|
|
ori_sequence_idx = [162, 127, 234, 93, 132, 58, 172, 136, 150, 149, 176, 148, 152, 377, 400, 378, 379, 365, 397, 288,
|
|
361, 323, 454, 356, 389, #
|
|
70, 63, 105, 66, 107, 55, 65, 52, 53, 46, #
|
|
336, 296, 334, 293, 300, 276, 283, 282, 295, 285, #
|
|
168, 6, 197, 195, 5, #
|
|
48, 115, 220, 45, 4, 275, 440, 344, 278, #
|
|
33, 246, 161, 160, 159, 158, 157, 173, 133, 155, 154, 153, 145, 144, 163, 7, #
|
|
362, 398, 384, 385, 386, 387, 388, 466, 263, 249, 390, 373, 374, 380, 381, 382, #
|
|
61, 185, 40, 39, 37, 0, 267, 269, 270, 409, 291, 375, 321, 405, 314, 17, 84, 181, 91, 146, #
|
|
78, 191, 80, 81, 82, 13, 312, 311, 310, 415, 308, 324, 318, 402, 317, 14, 87, 178, 88, 95]
|
|
|
|
# the following is the connections of landmarks for drawing sketch image
|
|
FACEMESH_LIPS = frozenset([(61, 146), (146, 91), (91, 181), (181, 84), (84, 17),
|
|
(17, 314), (314, 405), (405, 321), (321, 375),
|
|
(375, 291), (61, 185), (185, 40), (40, 39), (39, 37),
|
|
(37, 0), (0, 267),
|
|
(267, 269), (269, 270), (270, 409), (409, 291),
|
|
(78, 95), (95, 88), (88, 178), (178, 87), (87, 14),
|
|
(14, 317), (317, 402), (402, 318), (318, 324),
|
|
(324, 308), (78, 191), (191, 80), (80, 81), (81, 82),
|
|
(82, 13), (13, 312), (312, 311), (311, 310),
|
|
(310, 415), (415, 308)])
|
|
FACEMESH_LEFT_EYE = frozenset([(263, 249), (249, 390), (390, 373), (373, 374),
|
|
(374, 380), (380, 381), (381, 382), (382, 362),
|
|
(263, 466), (466, 388), (388, 387), (387, 386),
|
|
(386, 385), (385, 384), (384, 398), (398, 362)])
|
|
FACEMESH_LEFT_EYEBROW = frozenset([(276, 283), (283, 282), (282, 295),
|
|
(295, 285), (300, 293), (293, 334),
|
|
(334, 296), (296, 336)])
|
|
FACEMESH_RIGHT_EYE = frozenset([(33, 7), (7, 163), (163, 144), (144, 145),
|
|
(145, 153), (153, 154), (154, 155), (155, 133),
|
|
(33, 246), (246, 161), (161, 160), (160, 159),
|
|
(159, 158), (158, 157), (157, 173), (173, 133)])
|
|
FACEMESH_RIGHT_EYEBROW = frozenset([(46, 53), (53, 52), (52, 65), (65, 55),
|
|
(70, 63), (63, 105), (105, 66), (66, 107)])
|
|
FACEMESH_FACE_OVAL = frozenset([(389, 356), (356, 454),
|
|
(454, 323), (323, 361), (361, 288), (288, 397),
|
|
(397, 365), (365, 379), (379, 378), (378, 400),
|
|
(400, 377), (377, 152), (152, 148), (148, 176),
|
|
(176, 149), (149, 150), (150, 136), (136, 172),
|
|
(172, 58), (58, 132), (132, 93), (93, 234),
|
|
(234, 127), (127, 162)])
|
|
FACEMESH_NOSE = frozenset([(168, 6), (6, 197), (197, 195), (195, 5), (5, 4),
|
|
(4, 45), (45, 220), (220, 115), (115, 48),
|
|
(4, 275), (275, 440), (440, 344), (344, 278), ])
|
|
FACEMESH_CONNECTION = frozenset().union(*[
|
|
FACEMESH_LIPS, FACEMESH_LEFT_EYE, FACEMESH_LEFT_EYEBROW, FACEMESH_RIGHT_EYE,
|
|
FACEMESH_RIGHT_EYEBROW, FACEMESH_FACE_OVAL, FACEMESH_NOSE
|
|
])
|
|
|
|
full_face_landmark_sequence = [*list(range(0, 4)), *list(range(21, 25)), *list(range(25, 91)), #upper-half face
|
|
*list(range(4, 21)), # jaw
|
|
*list(range(91, 131))] # mouth
|
|
|
|
class LandmarkDict(dict):# Makes a dictionary that behave like an object to represent each landmark
|
|
def __init__(self, idx, x, y):
|
|
self['idx'] = idx
|
|
self['x'] = x
|
|
self['y'] = y
|
|
def __getattr__(self, name):
|
|
try:
|
|
return self[name]
|
|
except:
|
|
raise AttributeError(name)
|
|
def __setattr__(self, name, value):
|
|
self[name] = value
|
|
|
|
class IP_LAP_infer:
|
|
|
|
def __init__(self,T=5,Nl=15,ref_img_N=25,
|
|
img_size=128,mel_step_size=16,
|
|
face_det_batch_size=4,
|
|
checkpoints_path=""):
|
|
self.T = T
|
|
self.Nl = Nl
|
|
self.ref_img_N = ref_img_N
|
|
self.img_size = img_size
|
|
self.mel_step_size = mel_step_size
|
|
self.face_det_batch_size = face_det_batch_size
|
|
self.pads = [100,100,100,100]
|
|
self.device = "cuda" if cuda_malloc_supported() else "cpu"
|
|
|
|
self.mp_face_mesh = mp.solutions.face_mesh
|
|
self.drawing_spec = mp.solutions.drawing_utils.DrawingSpec(thickness=1, circle_radius=1)
|
|
self.lip_index = [0, 17]
|
|
self.all_landmarks_idx = self.summarize_landmark(FACEMESH_CONNECTION)
|
|
self.pose_landmark_idx = \
|
|
self.summarize_landmark(FACEMESH_NOSE.union(*[FACEMESH_RIGHT_EYEBROW, FACEMESH_RIGHT_EYE,
|
|
FACEMESH_LEFT_EYE, FACEMESH_LEFT_EYEBROW, ])).union(
|
|
[162, 127, 234, 93, 389, 356, 454, 323])
|
|
# pose landmarks are landmarks of the upper-half face(eyes,nose,cheek) that represents the pose information
|
|
|
|
self.content_landmark_idx = self.all_landmarks_idx - self.pose_landmark_idx
|
|
# content_landmark include landmarks of lip and jaw which are inferred from audio
|
|
|
|
|
|
landmark_gen_checkpoint_path = os.path.join(checkpoints_path, "landmarkgenerator_checkpoint.pth")
|
|
renderer_checkpoint_path = os.path.join(checkpoints_path, "renderer_checkpoint.pth")
|
|
self.landmark_generator_model = self.load_model(
|
|
model=Landmark_transformer(T=self.T, d_model=512, nlayers=4, nhead=4, dim_feedforward=1024, dropout=0.1),
|
|
path=landmark_gen_checkpoint_path)
|
|
self.renderer = self.load_model(model=Renderer(), path=renderer_checkpoint_path)
|
|
|
|
self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType.TWO_D, flip_input=False, device=self.device)
|
|
|
|
self.face_mask = FaceMask()
|
|
|
|
def __call__(self,video_file, audio_file, outfile):
|
|
temp_dir = os.path.join(os.path.dirname(outfile), NAME)
|
|
if not os.path.exists(temp_dir): os.makedirs(temp_dir, exist_ok=True)
|
|
##(1) Reading input video frames ###
|
|
print(f'[Step 1]Reading video frames ... from {video_file}', NAME)
|
|
if not os.path.isfile(video_file):
|
|
raise ValueError('the input video file does not exist')
|
|
elif video_file.split('.')[1] in ['jpg', 'png', 'jpeg']: #if input a single image for testing
|
|
ori_background_frames = [cv2.imread(video_file)]
|
|
else:
|
|
video_stream = cv2.VideoCapture(video_file)
|
|
fps = video_stream.get(cv2.CAP_PROP_FPS)
|
|
if fps != 25:
|
|
print(" input video fps:", fps,',converting to 25fps...')
|
|
tmp_file = os.path.join(temp_dir,'temp_25fps.mp4')
|
|
if os.path.exists(tmp_file): os.remove(tmp_file)
|
|
print(tmp_file)
|
|
command = 'ffmpeg -y -i ' + video_file + f' -r 25 {tmp_file}'
|
|
subprocess.call(command, shell=platform.system() != 'Windows',stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
|
|
video_file = tmp_file
|
|
video_stream.release()
|
|
video_stream = cv2.VideoCapture(video_file)
|
|
fps = video_stream.get(cv2.CAP_PROP_FPS)
|
|
assert fps == 25
|
|
|
|
ori_background_frames = [] #input videos frames (includes background as well as face)
|
|
frame_idx = 0
|
|
while 1:
|
|
still_reading, frame = video_stream.read()
|
|
if not still_reading:
|
|
video_stream.release()
|
|
break
|
|
ori_background_frames.append(frame)
|
|
frame_idx = frame_idx + 1
|
|
input_vid_len = len(ori_background_frames)
|
|
|
|
##(2) Extracting audio####
|
|
print(f'[Step 2]Extracting audio ... from {audio_file}', NAME)
|
|
if not audio_file.endswith('.wav'):
|
|
tmp_file = os.join.path(temp_dir,'temp.wav')
|
|
command = 'ffmpeg -y -i {} -strict -2 {}'.format(audio_file, tmp_file)
|
|
subprocess.call(command, shell=platform.system() != 'Windows', stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
|
|
audio_file = tmp_file
|
|
wav = audio.load_wav(audio_file, 16000)
|
|
mel = audio.melspectrogram(wav) # (H,W) extract mel-spectrum
|
|
##read audio mel into list###
|
|
mel_chunks = [] # each mel chunk correspond to 5 video frames, used to generate one video frame
|
|
mel_idx_multiplier = 80. / fps
|
|
mel_chunk_idx = 0
|
|
while 1:
|
|
start_idx = int(mel_chunk_idx * mel_idx_multiplier)
|
|
if start_idx + self.mel_step_size > len(mel[0]):
|
|
break
|
|
mel_chunks.append(mel[:, start_idx: start_idx + self.mel_step_size]) # mel for generate one video frame
|
|
mel_chunk_idx += 1
|
|
|
|
print('[Step 3]detect facial using face detection tool', NAME)
|
|
ori_face_frames, ori_face_coords = self.face_detect(ori_background_frames)
|
|
# print(len(ori_face_frames))
|
|
import gc; gc.collect(); torch.cuda.empty_cache()
|
|
|
|
##(3) detect facial landmarks using mediapipe tool
|
|
print('[Step 4]detect facial landmarks using mediapipe tool', NAME)
|
|
boxes = [] #bounding boxes of human face
|
|
lip_dists = [] #lip dists
|
|
#we define the lip dist(openness): distance between the midpoints of the upper lip and lower lip
|
|
face_crop_results = []
|
|
all_pose_landmarks, all_content_landmarks = [], [] #content landmarks include lip and jaw landmarks
|
|
with self.mp_face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1, refine_landmarks=True,
|
|
min_detection_confidence=0) as face_mesh:
|
|
# (1) get bounding boxes and lip dist
|
|
for frame_idx, full_frame in tqdm(enumerate(ori_face_frames),total=input_vid_len,
|
|
desc="get bounding boxes and lip dist"):
|
|
h, w = full_frame.shape[0], full_frame.shape[1]
|
|
results = face_mesh.process(cv2.cvtColor(full_frame, cv2.COLOR_BGR2RGB))
|
|
if not results.multi_face_landmarks:
|
|
raise NotImplementedError # not detect face
|
|
face_landmarks = results.multi_face_landmarks[0]
|
|
|
|
## calculate the lip dist
|
|
dx = face_landmarks.landmark[self.lip_index[0]].x - face_landmarks.landmark[self.lip_index[1]].x
|
|
dy = face_landmarks.landmark[self.lip_index[0]].y - face_landmarks.landmark[self.lip_index[1]].y
|
|
dist = np.linalg.norm((dx, dy))
|
|
lip_dists.append((frame_idx, dist))
|
|
|
|
# (1)get the marginal landmarks to crop face
|
|
x_min,x_max,y_min,y_max = 999,-999,999,-999
|
|
for idx, landmark in enumerate(face_landmarks.landmark):
|
|
if idx in self.all_landmarks_idx:
|
|
if landmark.x < x_min:
|
|
x_min = landmark.x
|
|
if landmark.x > x_max:
|
|
x_max = landmark.x
|
|
if landmark.y < y_min:
|
|
y_min = landmark.y
|
|
if landmark.y > y_max:
|
|
y_max = landmark.y
|
|
##########plus some pixel to the marginal region##########
|
|
#note:the landmarks coordinates returned by mediapipe range 0~1
|
|
plus_pixel = 25
|
|
x_min = max(x_min - plus_pixel / w, 0)
|
|
x_max = min(x_max + plus_pixel / w, 1)
|
|
|
|
y_min = max(y_min - plus_pixel / h, 0)
|
|
y_max = min(y_max + plus_pixel / h, 1)
|
|
y1, y2, x1, x2 = int(y_min * h), int(y_max * h), int(x_min * w), int(x_max * w)
|
|
boxes.append([y1, y2, x1, x2])
|
|
boxes = np.array(boxes)
|
|
|
|
# (2)croppd face
|
|
face_crop_results = [[image[y1:y2, x1:x2], (y1, y2, x1, x2)] \
|
|
for image, (y1, y2, x1, x2) in zip(ori_face_frames, boxes)]
|
|
|
|
# (3)detect facial landmarks
|
|
for frame_idx, full_frame in tqdm(enumerate(ori_face_frames),total=input_vid_len,
|
|
desc="detect facial landmarks"):
|
|
h, w = full_frame.shape[0], full_frame.shape[1]
|
|
results = face_mesh.process(cv2.cvtColor(full_frame, cv2.COLOR_BGR2RGB))
|
|
if not results.multi_face_landmarks:
|
|
raise ValueError("not detect face in some frame!") # not detect
|
|
face_landmarks = results.multi_face_landmarks[0]
|
|
|
|
|
|
|
|
pose_landmarks, content_landmarks = [], []
|
|
for idx, landmark in enumerate(face_landmarks.landmark):
|
|
if idx in self.pose_landmark_idx:
|
|
pose_landmarks.append((idx, w * landmark.x, h * landmark.y))
|
|
if idx in self.content_landmark_idx:
|
|
content_landmarks.append((idx, w * landmark.x, h * landmark.y))
|
|
|
|
# normalize landmarks to 0~1
|
|
y_min, y_max, x_min, x_max = face_crop_results[frame_idx][1] #bounding boxes
|
|
pose_landmarks = [ \
|
|
[idx, (x - x_min) / (x_max - x_min), (y - y_min) / (y_max - y_min)] for idx, x, y in pose_landmarks]
|
|
content_landmarks = [ \
|
|
[idx, (x - x_min) / (x_max - x_min), (y - y_min) / (y_max - y_min)] for idx, x, y in content_landmarks]
|
|
all_pose_landmarks.append(pose_landmarks)
|
|
all_content_landmarks.append(content_landmarks)
|
|
|
|
all_pose_landmarks = self.get_smoothened_landmarks(all_pose_landmarks, windows_T=1)
|
|
all_content_landmarks=self.get_smoothened_landmarks(all_content_landmarks,windows_T=1)
|
|
|
|
##randomly select N_l reference landmarks for landmark transformer##
|
|
print("randomly select N_l reference landmarks for landmark transformer", NAME)
|
|
dists_sorted = sorted(lip_dists, key=lambda x: x[1])
|
|
lip_dist_idx = np.asarray([idx for idx, dist in dists_sorted]) #the frame idxs sorted by lip openness
|
|
|
|
Nl_idxs = [lip_dist_idx[int(i)] for i in torch.linspace(0, input_vid_len - 1, steps=self.Nl)]
|
|
Nl_pose_landmarks, Nl_content_landmarks = [], [] #Nl_pose + Nl_content=Nl reference landmarks
|
|
for reference_idx in Nl_idxs:
|
|
frame_pose_landmarks = all_pose_landmarks[reference_idx]
|
|
frame_content_landmarks = all_content_landmarks[reference_idx]
|
|
Nl_pose_landmarks.append(frame_pose_landmarks)
|
|
Nl_content_landmarks.append(frame_content_landmarks)
|
|
|
|
Nl_pose = torch.zeros((self.Nl, 2, 74)) # 74 landmark
|
|
Nl_content = torch.zeros((self.Nl, 2, 57)) # 57 landmark
|
|
for idx in range(self.Nl):
|
|
#arrange the landmark in a certain order, since the landmark index returned by mediapipe is is chaotic
|
|
Nl_pose_landmarks[idx] = sorted(Nl_pose_landmarks[idx],
|
|
key=lambda land_tuple: ori_sequence_idx.index(land_tuple[0]))
|
|
Nl_content_landmarks[idx] = sorted(Nl_content_landmarks[idx],
|
|
key=lambda land_tuple: ori_sequence_idx.index(land_tuple[0]))
|
|
|
|
Nl_pose[idx, 0, :] = torch.FloatTensor(
|
|
[Nl_pose_landmarks[idx][i][1] for i in range(len(Nl_pose_landmarks[idx]))]) # x
|
|
Nl_pose[idx, 1, :] = torch.FloatTensor(
|
|
[Nl_pose_landmarks[idx][i][2] for i in range(len(Nl_pose_landmarks[idx]))]) # y
|
|
Nl_content[idx, 0, :] = torch.FloatTensor(
|
|
[Nl_content_landmarks[idx][i][1] for i in range(len(Nl_content_landmarks[idx]))]) # x
|
|
Nl_content[idx, 1, :] = torch.FloatTensor(
|
|
[Nl_content_landmarks[idx][i][2] for i in range(len(Nl_content_landmarks[idx]))]) # y
|
|
Nl_content = Nl_content.unsqueeze(0) # (1,Nl, 2, 57)
|
|
Nl_pose = Nl_pose.unsqueeze(0) # (1,Nl,2,74)
|
|
|
|
|
|
##select reference images and draw sketches for rendering according to lip openness##
|
|
print("select reference images and draw sketches for rendering according to lip openness", NAME)
|
|
ref_img_idx = [int(lip_dist_idx[int(i)]) for i in torch.linspace(0, input_vid_len - 1, steps=self.ref_img_N)]
|
|
ref_imgs = [face_crop_results[idx][0] for idx in ref_img_idx]
|
|
## (N,H,W,3)
|
|
ref_img_pose_landmarks, ref_img_content_landmarks = [], []
|
|
for idx in ref_img_idx:
|
|
ref_img_pose_landmarks.append(all_pose_landmarks[idx])
|
|
ref_img_content_landmarks.append(all_content_landmarks[idx])
|
|
|
|
ref_img_pose = torch.zeros((self.ref_img_N, 2, 74)) # 74 landmark
|
|
ref_img_content = torch.zeros((self.ref_img_N, 2, 57)) # 57 landmark
|
|
|
|
for idx in range(self.ref_img_N):
|
|
ref_img_pose_landmarks[idx] = sorted(ref_img_pose_landmarks[idx],
|
|
key=lambda land_tuple: ori_sequence_idx.index(land_tuple[0]))
|
|
ref_img_content_landmarks[idx] = sorted(ref_img_content_landmarks[idx],
|
|
key=lambda land_tuple: ori_sequence_idx.index(land_tuple[0]))
|
|
ref_img_pose[idx, 0, :] = torch.FloatTensor(
|
|
[ref_img_pose_landmarks[idx][i][1] for i in range(len(ref_img_pose_landmarks[idx]))]) # x
|
|
ref_img_pose[idx, 1, :] = torch.FloatTensor(
|
|
[ref_img_pose_landmarks[idx][i][2] for i in range(len(ref_img_pose_landmarks[idx]))]) # y
|
|
|
|
ref_img_content[idx, 0, :] = torch.FloatTensor(
|
|
[ref_img_content_landmarks[idx][i][1] for i in range(len(ref_img_content_landmarks[idx]))]) # x
|
|
ref_img_content[idx, 1, :] = torch.FloatTensor(
|
|
[ref_img_content_landmarks[idx][i][2] for i in range(len(ref_img_content_landmarks[idx]))]) # y
|
|
|
|
ref_img_full_face_landmarks = torch.cat([ref_img_pose, ref_img_content], dim=2).cpu().numpy() # (N,2,131)
|
|
ref_img_sketches = []
|
|
for frame_idx in range(ref_img_full_face_landmarks.shape[0]): # N
|
|
full_landmarks = ref_img_full_face_landmarks[frame_idx] # (2,131)
|
|
h, w = ref_imgs[frame_idx].shape[0], ref_imgs[frame_idx].shape[1]
|
|
drawn_sketech = np.zeros((int(h * self.img_size / min(h, w)), int(w * self.img_size / min(h, w)), 3))
|
|
mediapipe_format_landmarks = [LandmarkDict(ori_sequence_idx[full_face_landmark_sequence[idx]], full_landmarks[0, idx],
|
|
full_landmarks[1, idx]) for idx in range(full_landmarks.shape[1])]
|
|
drawn_sketech = draw_landmarks(drawn_sketech, mediapipe_format_landmarks, connections=FACEMESH_CONNECTION,
|
|
connection_drawing_spec=self.drawing_spec)
|
|
drawn_sketech = cv2.resize(drawn_sketech, (self.img_size, self.img_size)) # (128, 128, 3)
|
|
ref_img_sketches.append(drawn_sketech)
|
|
ref_img_sketches = torch.FloatTensor(np.asarray(ref_img_sketches) / 255.0).cuda().unsqueeze(0).permute(0, 1, 4, 2, 3)
|
|
# (1,N, 3, 128, 128)
|
|
ref_imgs = [cv2.resize(face.copy(), (self.img_size, self.img_size)) for face in ref_imgs]
|
|
ref_imgs = torch.FloatTensor(np.asarray(ref_imgs) / 255.0).unsqueeze(0).permute(0, 1, 4, 2, 3).cuda()
|
|
# (1,N,3,H,W)
|
|
|
|
##prepare output video strame##
|
|
frame_h, frame_w = ori_background_frames[0].shape[:-1]
|
|
'''
|
|
out_stream = cv2.VideoWriter('{}/result.avi'.format(temp_dir), cv2.VideoWriter_fourcc(*'DIVX'), fps,
|
|
(frame_w, frame_h)) # +frame_h*3
|
|
'''
|
|
out_stream = cv2.VideoWriter(outfile, cv2.VideoWriter_fourcc(*'mp4v'), fps,
|
|
(frame_w, frame_h)) # +frame_h*3
|
|
|
|
##generate final face image and output video##
|
|
input_mel_chunks_len = len(mel_chunks)
|
|
input_frame_sequence = torch.arange(input_vid_len).tolist()
|
|
#the input template video may be shorter than audio
|
|
#in this case we repeat the input template video as following
|
|
num_of_repeat=input_mel_chunks_len//input_vid_len+1
|
|
input_frame_sequence = input_frame_sequence + list(reversed(input_frame_sequence))
|
|
input_frame_sequence=input_frame_sequence*((num_of_repeat+1)//2)
|
|
file_num = 0
|
|
for batch_idx, batch_start_idx in tqdm(enumerate(range(0, input_mel_chunks_len-2, 1)),
|
|
total=len(range(0, input_mel_chunks_len-2, 1)), desc="[IP_LAP] [Step 5]Lipsync..."):
|
|
T_input_frame, T_ori_face_coordinates = [], []
|
|
#note: input_frame include background as well as face
|
|
T_mel_batch, T_crop_face,T_pose_landmarks = [], [],[]
|
|
|
|
A_input_frame, A_ori_face_coordinates = [], []
|
|
|
|
# (1) for each batch of T frame, generate corresponding landmarks using landmark generator
|
|
for mel_chunk_idx in range(batch_start_idx, batch_start_idx + self.T): # for each T frame
|
|
# 1 input audio
|
|
T_mel_batch.append(mel_chunks[max(0, mel_chunk_idx - 2)])
|
|
|
|
# 2.input face
|
|
input_frame_idx = int(input_frame_sequence[mel_chunk_idx])
|
|
face, coords = face_crop_results[input_frame_idx]
|
|
T_crop_face.append(face)
|
|
T_ori_face_coordinates.append((face, coords)) ##input face
|
|
# 3.pose landmarks
|
|
T_pose_landmarks.append(all_pose_landmarks[input_frame_idx])
|
|
# 3.face background
|
|
T_input_frame.append(ori_face_frames[input_frame_idx].copy())
|
|
# 4.frame background
|
|
A_ori_face_coordinates.append(ori_face_coords[input_frame_idx])
|
|
A_input_frame.append(ori_background_frames[input_frame_idx].copy())
|
|
|
|
T_mels = torch.FloatTensor(np.asarray(T_mel_batch)).unsqueeze(1).unsqueeze(0) # 1,T,1,h,w
|
|
#prepare pose landmarks
|
|
T_pose = torch.zeros((self.T, 2, 74)) # 74 landmark
|
|
for idx in range(self.T):
|
|
T_pose_landmarks[idx] = sorted(T_pose_landmarks[idx],
|
|
key=lambda land_tuple: ori_sequence_idx.index(land_tuple[0]))
|
|
T_pose[idx, 0, :] = torch.FloatTensor(
|
|
[T_pose_landmarks[idx][i][1] for i in range(len(T_pose_landmarks[idx]))]) # x
|
|
T_pose[idx, 1, :] = torch.FloatTensor(
|
|
[T_pose_landmarks[idx][i][2] for i in range(len(T_pose_landmarks[idx]))]) # y
|
|
T_pose = T_pose.unsqueeze(0) # (1,T, 2,74)
|
|
|
|
#landmark generator inference
|
|
Nl_pose, Nl_content = Nl_pose.cuda(), Nl_content.cuda() # (Nl,2,74) (Nl,2,57)
|
|
T_mels, T_pose = T_mels.cuda(), T_pose.cuda()
|
|
with torch.no_grad(): # require (1,T,1,hv,wv)(1,T,2,74)(1,T,2,57)
|
|
predict_content = self.landmark_generator_model(T_mels, T_pose, Nl_pose, Nl_content) # (1*T,2,57)
|
|
T_pose = torch.cat([T_pose[i] for i in range(T_pose.size(0))], dim=0) # (1*T,2,74)
|
|
T_predict_full_landmarks = torch.cat([T_pose, predict_content], dim=2).cpu().numpy() # (1*T,2,131)
|
|
|
|
#1.draw target sketch
|
|
T_target_sketches = []
|
|
for frame_idx in range(self.T):
|
|
full_landmarks = T_predict_full_landmarks[frame_idx] # (2,131)
|
|
h, w = T_crop_face[frame_idx].shape[0], T_crop_face[frame_idx].shape[1]
|
|
drawn_sketech = np.zeros((int(h * self.img_size / min(h, w)), int(w * self.img_size / min(h, w)), 3))
|
|
mediapipe_format_landmarks = [LandmarkDict(ori_sequence_idx[full_face_landmark_sequence[idx]]
|
|
, full_landmarks[0, idx], full_landmarks[1, idx]) for idx in
|
|
range(full_landmarks.shape[1])]
|
|
drawn_sketech = draw_landmarks(drawn_sketech, mediapipe_format_landmarks, connections=FACEMESH_CONNECTION,
|
|
connection_drawing_spec=self.drawing_spec)
|
|
drawn_sketech = cv2.resize(drawn_sketech, (self.img_size, self.img_size)) # (128, 128, 3)
|
|
if frame_idx == 2:
|
|
show_sketch = cv2.resize(drawn_sketech, (frame_w, frame_h)).astype(np.uint8)
|
|
T_target_sketches.append(torch.FloatTensor(drawn_sketech) / 255)
|
|
T_target_sketches = torch.stack(T_target_sketches, dim=0).permute(0, 3, 1, 2) # (T,3,128, 128)
|
|
target_sketches = T_target_sketches.unsqueeze(0).cuda() # (1,T,3,128, 128)
|
|
|
|
# 2.lower-half masked face
|
|
ori_face_img = torch.FloatTensor(cv2.resize(T_crop_face[2], (self.img_size, self.img_size)) / 255).permute(2, 0, 1).unsqueeze(
|
|
0).unsqueeze(0).cuda() #(1,1,3,H, W)
|
|
|
|
# 3. render the full face
|
|
# require (1,1,3,H,W) (1,T,3,H,W) (1,N,3,H,W) (1,N,3,H,W) (1,1,1,h,w)
|
|
# return (1,3,H,W)
|
|
with torch.no_grad():
|
|
generated_face, _, _, _ = self.renderer(ori_face_img, target_sketches, ref_imgs, ref_img_sketches,
|
|
T_mels[:, 2].unsqueeze(0)) # T=1
|
|
gen_face = (generated_face.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8) # (H,W,3)
|
|
|
|
# 4. paste each generated face
|
|
y1, y2, x1, x2 = T_ori_face_coordinates[2][1] # coordinates of face bounding box
|
|
original_background = T_input_frame[2].copy()
|
|
T_input_frame[2][y1:y2, x1:x2] = cv2.resize(gen_face,(x2 - x1, y2 - y1)) #resize and paste generated face
|
|
# 5. post-process
|
|
full_face = self.merge_face_contour_only(original_background, T_input_frame[2], T_ori_face_coordinates[2][1],self.fa) #(H,W,3)
|
|
# 6.output
|
|
# full = np.concatenate([show_sketch, full], axis=1)
|
|
# print(f"full_face.shape{full_face.shape}")
|
|
ori_x1, ori_y1, ori_x2, ori_y2 = A_ori_face_coordinates[2]
|
|
# print(ori_face_coords[file_num])
|
|
full_frame = A_input_frame[2]
|
|
#full_mask = np.zeros_like(full_frame)
|
|
# print(full_frame.shape)
|
|
if ori_x1 != -1:
|
|
p = cv2.resize(full_face.astype(np.uint8), (ori_x2 - ori_x1, ori_y2 - ori_y1))
|
|
# print(p.shape)
|
|
full_frame[ori_y1:ori_y2, ori_x1:ori_x2] = p
|
|
# height, width = full_frame.shape[:2]
|
|
# img = self.Laplacian_Pyramid_Blending_with_mask(full_frame, ori_background_frames[file_num], full_mask[:, :, 0], 6)
|
|
# pp = np.uint8(cv2.resize(np.clip(img, 0 ,255), (width, height)))
|
|
mask = self.face_mask(p)
|
|
full_frame[ori_y1:ori_y2, ori_x1:ori_x2] = full_frame[ori_y1:ori_y2, ori_x1:ori_x2] * (1 - mask[..., None]) + p * mask[..., None]
|
|
full = full_frame.copy()
|
|
|
|
out_stream.write(full)
|
|
|
|
try:
|
|
# cv2.imwrite(temp_frame_paths[batch_idx+2],full)
|
|
file_num += 1
|
|
except:
|
|
pass
|
|
|
|
if batch_idx == 0:
|
|
out_stream.write(full)
|
|
out_stream.write(full)
|
|
# cv2.imwrite(temp_frame_paths[batch_idx],full)
|
|
# cv2.imwrite(temp_frame_paths[batch_idx+1],full)
|
|
|
|
out_stream.release()
|
|
# command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(voice_file, '{}/result.avi'.format(temp_dir), outfile)
|
|
# subprocess.call(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
|
|
print(f"succeed output results to:{outfile}", NAME)
|
|
|
|
|
|
def face_detect(self, images):
|
|
|
|
batch_size = self.face_det_batch_size
|
|
|
|
while 1:
|
|
predictions = []
|
|
try:
|
|
for i in tqdm(range(0, len(images), batch_size)):
|
|
imgs = np.array(images[i:i + batch_size])
|
|
imgs_numpy = imgs.transpose(0, 3, 1, 2)
|
|
image_batch = torch.from_numpy(imgs_numpy.copy())
|
|
_, _, bboxes =self.fa.get_landmarks_from_batch(image_batch,return_bboxes=True)
|
|
predictions.extend(bboxes)
|
|
except RuntimeError:
|
|
if batch_size == 1:
|
|
raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument')
|
|
batch_size //= 2
|
|
print('Recovering from OOM error; New batch size: {}'.format(batch_size))
|
|
continue
|
|
break
|
|
|
|
results = []
|
|
pady1, pady2, padx1, padx2 = self.pads
|
|
for rect, image in zip(predictions, images):
|
|
|
|
if rect is None:
|
|
# cv2.imwrite('temp/faulty_frame.jpg', image) # check this frame where the face was not detected.
|
|
# results.append([-1,-1,-1,-1])
|
|
raise ValueError('Face not detected! Ensure the video contains a face in all the frames.')
|
|
else:
|
|
rect = rect[0]
|
|
rect = np.clip(rect, 0, None)
|
|
x1_0, y1_0, x2_0, y2_0 = map(int, rect[:-1])
|
|
|
|
y1 = max(0, y1_0 - pady1)
|
|
y2 = min(image.shape[0], y2_0 + pady2)
|
|
x1 = max(0, x1_0 - padx1)
|
|
x2 = min(image.shape[1], x2_0 + padx2)
|
|
|
|
results.append([x1, y1, x2, y2])
|
|
|
|
boxes = np.array(results)
|
|
|
|
faces = [image[y1: y2, x1:x2] for image, (x1, y1, x2, y2) in zip(images, boxes)]
|
|
return faces, boxes
|
|
|
|
|
|
def merge_face_contour_only(self,src_frame, generated_frame, face_region_coord, fa): #function used in post-process
|
|
"""Merge the face from generated_frame into src_frame
|
|
"""
|
|
input_img = src_frame
|
|
y1, y2, x1, x2 = 0, 0, 0, 0
|
|
if face_region_coord is not None:
|
|
y1, y2, x1, x2 = face_region_coord
|
|
input_img = src_frame[y1:y2, x1:x2]
|
|
### 1) Detect the facial landmarks
|
|
try:
|
|
preds = fa.get_landmarks(input_img)[0] # 68x2
|
|
except:
|
|
preds = np.int64(-1 * np.ones((68,2)))
|
|
if face_region_coord is not None:
|
|
preds += np.array([x1, y1])
|
|
lm_pts = preds.astype(int)
|
|
contour_idx = list(range(0, 17)) + list(range(17, 27))[::-1]
|
|
contour_pts = lm_pts[contour_idx]
|
|
### 2) Make the landmark region mark image
|
|
mask_img = np.zeros((src_frame.shape[0], src_frame.shape[1], 1), np.uint8)
|
|
cv2.fillConvexPoly(mask_img, contour_pts, 255)
|
|
### 3) Do swap
|
|
img = self.swap_masked_region(src_frame, generated_frame, mask=mask_img)
|
|
return img
|
|
|
|
def swap_masked_region(self,target_img, src_img, mask): #function used in post-process
|
|
"""From src_img crop masked region to replace corresponding masked region
|
|
in target_img
|
|
""" # swap_masked_region(src_frame, generated_frame, mask=mask_img)
|
|
mask_img = cv2.GaussianBlur(mask, (21, 21), 11)
|
|
mask1 = mask_img / 255
|
|
mask1 = np.tile(np.expand_dims(mask1, axis=2), (1, 1, 3))
|
|
img = src_img * mask1 + target_img * (1 - mask1)
|
|
return img.astype(np.uint8)
|
|
|
|
# smooth landmarks
|
|
def get_smoothened_landmarks(self,all_landmarks, windows_T=1):
|
|
for i in range(len(all_landmarks)): # frame i
|
|
if i + windows_T > len(all_landmarks):
|
|
window = all_landmarks[len(all_landmarks) - windows_T:]
|
|
else:
|
|
window = all_landmarks[i: i + windows_T]
|
|
#####
|
|
for j in range(len(all_landmarks[i])): # landmark j
|
|
all_landmarks[i][j][1] = np.mean([frame_landmarks[j][1] for frame_landmarks in window]) # x
|
|
all_landmarks[i][j][2] = np.mean([frame_landmarks[j][2] for frame_landmarks in window]) # y
|
|
return all_landmarks
|
|
|
|
def load_model(self, model, path):
|
|
print("Load checkpoint from: {}".format(path))
|
|
checkpoint = self._load(path)
|
|
s = checkpoint["state_dict"]
|
|
new_s = {}
|
|
for k, v in s.items():
|
|
if k[:6] == 'module':
|
|
new_k=k.replace('module.', '', 1)
|
|
else:
|
|
new_k =k
|
|
new_s[new_k] = v
|
|
model.load_state_dict(new_s)
|
|
model = model.to(self.device)
|
|
return model.eval()
|
|
|
|
def _load(self,checkpoint_path):
|
|
if self.device == 'cuda':
|
|
checkpoint = torch.load(checkpoint_path)
|
|
else:
|
|
checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
|
|
return checkpoint
|
|
|
|
def summarize_landmark(self, edge_set): # summarize all ficial landmarks used to construct edge
|
|
landmarks = set()
|
|
for a, b in edge_set:
|
|
landmarks.add(a)
|
|
landmarks.add(b)
|
|
return landmarks
|