# Prediction interface for Cog ⚙️
# https://github.com/replicate/cog/blob/main/docs/python.md
import os
import sys
import argparse
import subprocess
import numpy as np
from tqdm import tqdm
from PIL import Image
from scipy.io import loadmat
import torch
import cv2
from cog import BasePredictor, Input, Path
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,
)
class Predictor(BasePredictor):
def setup(self) -> None:
"""Load the model into memory to make running multiple predictions efficient"""
self.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="cuda",
)
self.restorer = GFPGANer(
model_path="checkpoints/GFPGANv1.3.pth",
upscale=1,
arch="clean",
channel_multiplier=2,
bg_upsampler=None,
)
self.croper = Croper("checkpoints/shape_predictor_68_face_landmarks.dat")
self.kp_extractor = KeypointExtractor()
face3d_net_path = "checkpoints/face3d_pretrain_epoch_20.pth"
self.net_recon = load_face3d_net(face3d_net_path, "cuda")
self.lm3d_std = load_lm3d("checkpoints/BFM")
def predict(
self,
face: Path = Input(description="Input video file of a talking-head."),
input_audio: Path = Input(description="Input audio file."),
) -> Path:
"""Run a single prediction on the model"""
device = "cuda"
args = argparse.Namespace(
DNet_path="checkpoints/DNet.pt",
LNet_path="checkpoints/LNet.pth",
ENet_path="checkpoints/ENet.pth",
face3d_net_path="checkpoints/face3d_pretrain_epoch_20.pth",
face=str(face),
audio=str(input_audio),
exp_img="neutral",
outfile=None,
fps=25,
pads=[0, 20, 0, 0],
face_det_batch_size=4,
LNet_batch_size=16,
img_size=384,
crop=[0, -1, 0, -1],
box=[-1, -1, -1, -1],
nosmooth=False,
static=False,
up_face="original",
one_shot=False,
without_rl1=False,
tmp_dir="temp",
re_preprocess=False,
)
base_name = args.face.split("/")[-1]
if args.face.split(".")[1] in ["jpg", "png", "jpeg"]:
full_frames = [cv2.imread(args.face)]
args.static = True
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)
full_frames_RGB = [
cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in full_frames
]
full_frames_RGB, crop, quad = self.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.")
lm = self.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
):
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 = (self.lm3d_std[:, :2] + 1) / 2.0
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, self.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.0, dtype=torch.float32)
.permute(2, 0, 1)
.to(device)
.unsqueeze(0)
)
with torch.no_grad():
coeffs = split_coeff(self.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 == "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] Stabilize 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.0
* 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 stabilized 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.0 / 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, _, _ = self.enhancer.process(
img, img, face_enhance=True, possion_blending=False
)
imgs_enhanced.append(pred)
gen = datagen(
imgs_enhanced.copy(), mel_chunks, full_frames, args, (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.0
) # 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 + 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.0
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 = self.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 = self.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.0
)
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 = self.enhancer.process(
pp, xf, bbox=c, face_enhance=False, possion_blending=True
)
out.write(pp)
out.release()
output_file = "/tmp/output.mp4"
command = "ffmpeg -loglevel error -y -i {} -i {} -strict -2 -q:v 1 {}".format(
args.audio, "temp/{}/result.mp4".format(args.tmp_dir), output_file
)
subprocess.call(command, shell=True)
return Path(output_file)
# frames:256x256, full_frames: original size
def datagen(frames, mels, full_frames, args, 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.0
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.0
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