import time import torch import onnx import cv2 import onnxruntime import numpy as np from tqdm import tqdm from onnx import numpy_helper from skimage import transform as trans import torchvision.transforms.functional as F from utils import make_white_image, laplacian_blending arcface_dst = np.array( [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], [41.5493, 92.3655], [70.7299, 92.2041]], dtype=np.float32) def estimate_norm(lmk, image_size=112, mode='arcface'): assert lmk.shape == (5, 2) assert image_size % 112 == 0 or image_size % 128 == 0 if image_size % 112 == 0: ratio = float(image_size) / 112.0 diff_x = 0 else: ratio = float(image_size) / 128.0 diff_x = 8.0 * ratio dst = arcface_dst * ratio dst[:, 0] += diff_x tform = trans.SimilarityTransform() tform.estimate(lmk, dst) M = tform.params[0:2, :] return M def norm_crop2(img, landmark, image_size=112, mode='arcface'): M = estimate_norm(landmark, image_size, mode) warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0) return warped, M class Inswapper(): def __init__(self, model_file=None, batch_size=32, providers=['CPUExecutionProvider']): self.model_file = model_file self.batch_size = batch_size model = onnx.load(self.model_file) graph = model.graph self.emap = numpy_helper.to_array(graph.initializer[-1]) self.input_mean = 0.0 self.input_std = 255.0 self.session_options = onnxruntime.SessionOptions() self.session = onnxruntime.InferenceSession(self.model_file, sess_options=self.session_options, providers=providers) inputs = self.session.get_inputs() self.input_names = [inp.name for inp in inputs] outputs = self.session.get_outputs() self.output_names = [out.name for out in outputs] assert len(self.output_names) == 1 self.output_shape = outputs[0].shape input_cfg = inputs[0] input_shape = input_cfg.shape self.input_shape = input_shape self.input_size = tuple(input_shape[2:4][::-1]) def forward(self, imgs, latents): preds = [] for img, latent in zip(imgs, latents): img = (img - self.input_mean) / self.input_std pred = self.session.run(self.output_names, {self.input_names[0]: img, self.input_names[1]: latent})[0] preds.append(pred) def get(self, imgs, target_faces, source_faces): imgs = list(imgs) preds = [None] * len(imgs) aimgs = [None] * len(imgs) matrs = [None] * len(imgs) for idx, (img, target_face, source_face) in enumerate(zip(imgs, target_faces, source_faces)): aimg, M, blob, latent = self.prepare_data(img, target_face, source_face) aimgs[idx] = aimg matrs[idx] = M pred = self.session.run(self.output_names, {self.input_names[0]: blob, self.input_names[1]: latent})[0] pred = pred.transpose((0, 2, 3, 1))[0] pred = np.clip(255 * pred, 0, 255).astype(np.uint8)[:, :, ::-1] preds[idx] = pred return (preds, aimgs, matrs) def prepare_data(self, img, target_face, source_face): if isinstance(img, str): img = cv2.imread(img) aimg, M = norm_crop2(img, target_face.kps, self.input_size[0]) blob = cv2.dnn.blobFromImage(aimg, 1.0 / self.input_std, self.input_size, (self.input_mean, self.input_mean, self.input_mean), swapRB=True) latent = source_face.normed_embedding.reshape((1, -1)) latent = np.dot(latent, self.emap) latent /= np.linalg.norm(latent) return (aimg, M, blob, latent) def batch_forward(self, img_list, target_f_list, source_f_list): num_samples = len(img_list) num_batches = (num_samples + self.batch_size - 1) // self.batch_size preds = [] aimgs = [] matrs = [] for i in tqdm(range(num_batches), desc="Swapping face"): start_idx = i * self.batch_size end_idx = min((i + 1) * self.batch_size, num_samples) batch_img = img_list[start_idx:end_idx] batch_target_f = target_f_list[start_idx:end_idx] batch_source_f = source_f_list[start_idx:end_idx] batch_pred, batch_aimg, batch_matr = self.get(batch_img, batch_target_f, batch_source_f) preds.extend(batch_pred) aimgs.extend(batch_aimg) matrs.extend(batch_matr) return (preds, aimgs, matrs) def paste_to_whole(bgr_fake, aimg, M, whole_img, laplacian_blend=True, crop_mask=(0,0,0,0)): IM = cv2.invertAffineTransform(M) img_white = make_white_image(aimg.shape[:2], crop=crop_mask, white_value=255) bgr_fake = cv2.warpAffine(bgr_fake, IM, (whole_img.shape[1], whole_img.shape[0]), borderValue=0.0) img_white = cv2.warpAffine(img_white, IM, (whole_img.shape[1], whole_img.shape[0]), borderValue=0.0) img_white[img_white > 20] = 255 img_mask = img_white mask_h_inds, mask_w_inds = np.where(img_mask == 255) mask_size = int(np.sqrt(np.ptp(mask_h_inds) * np.ptp(mask_w_inds))) k = max(mask_size // 10, 10) img_mask = cv2.erode(img_mask, np.ones((k, k), np.uint8), iterations=1) k = max(mask_size // 20, 5) kernel_size = (k, k) blur_size = tuple(2 * i + 1 for i in kernel_size) img_mask = cv2.GaussianBlur(img_mask, blur_size, 0) / 255 img_mask = np.tile(np.expand_dims(img_mask, axis=-1), (1, 1, 3)) if laplacian_blend: bgr_fake = laplacian_blending(bgr_fake.astype("float32").clip(0,255), whole_img.astype("float32").clip(0,255), img_mask.clip(0,1)) bgr_fake = bgr_fake.astype("float32") fake_merged = img_mask * bgr_fake + (1 - img_mask) * whole_img.astype(np.float32) return fake_merged.astype("uint8") def place_foreground_on_background(foreground, background, matrix): matrix = cv2.invertAffineTransform(matrix) mask = np.ones(foreground.shape, dtype="float32") foreground = cv2.warpAffine(foreground, matrix, (background.shape[1], background.shape[0]), borderValue=0.0) mask = cv2.warpAffine(mask, matrix, (background.shape[1], background.shape[0]), borderValue=0.0) composite_image = mask * foreground + (1 - mask) * background return composite_image