import os
import cv2
import gradio
import numpy as np
from huggingface_hub import Repository
from scipy.ndimage import gaussian_filter
from tensorflow.keras.models import load_model
from tensorflow_addons.layers import InstanceNormalization
from networks.layers import AdaIN, AdaptiveAttention
from options.swap_options import SwapOptions
from utils.utils import (estimate_norm, get_lm, inverse_estimate_norm,
norm_crop, transform_landmark_points)
# .
# token = os.environ['model_fetch']
opt = SwapOptions().parse()
retina_repo = Repository(local_dir="retina_model", clone_from="felixrosberg/RetinaFace")
from retina_model.models import *
RetinaFace = load_model("retina_model/retinaface_res50.h5",
custom_objects={"FPN": FPN,
"SSH": SSH,
"BboxHead": BboxHead,
"LandmarkHead": LandmarkHead,
"ClassHead": ClassHead})
arc_repo = Repository(local_dir="arcface_model", clone_from="felixrosberg/ArcFace")
ArcFace = load_model("arcface_model/arc_res50.h5")
ArcFaceE = load_model("arcface_model/arc_res50e.h5")
g_repo = Repository(local_dir="g_model_c_hq", clone_from="felixrosberg/FaceDancer")
G = load_model("g_model_c_hq/FaceDancer_config_c_HQ.h5", custom_objects={"AdaIN": AdaIN,
"AdaptiveAttention": AdaptiveAttention,
"InstanceNormalization": InstanceNormalization})
# r_repo = Repository(local_dir="reconstruction_attack", clone_from="felixrosberg/reconstruction_attack",
# private=True, use_auth_token=token)
# R = load_model("reconstruction_attack/reconstructor_42.h5", custom_objects={"AdaIN": AdaIN,
# "AdaptiveAttention": AdaptiveAttention,
# "InstanceNormalization": InstanceNormalization})
# permuter_repo = Repository(local_dir="identity_permuter", clone_from="felixrosberg/identitypermuter",
# private=True, use_auth_token=token, git_user="felixrosberg")
# from identity_permuter.id_permuter import identity_permuter
# IDP = identity_permuter(emb_size=32, min_arg=False)
# IDP.load_weights("identity_permuter/id_permuter.h5")
blend_mask_base = np.zeros(shape=(256, 256, 1))
blend_mask_base[80:244, 32:224] = 1
blend_mask_base = gaussian_filter(blend_mask_base, sigma=7)
def run_inference(target, source, slider, adv_slider, settings):
try:
source = np.array(source)
target = np.array(target)
# Prepare to load video
if "anonymize" not in settings:
source_a = RetinaFace(np.expand_dims(source, axis=0)).numpy()[0]
source_h, source_w, _ = source.shape
source_lm = get_lm(source_a, source_w, source_h)
source_aligned = norm_crop(source, source_lm, image_size=256)
source_z = ArcFace.predict(np.expand_dims(tf.image.resize(source_aligned, [112, 112]) / 255.0, axis=0))
else:
source_z = None
# read frame
im = target
im_h, im_w, _ = im.shape
im_shape = (im_w, im_h)
detection_scale = im_w // 640 if im_w > 640 else 1
faces = RetinaFace(np.expand_dims(cv2.resize(im,
(im_w // detection_scale,
im_h // detection_scale)), axis=0)).numpy()
total_img = im / 255.0
for annotation in faces:
lm_align = np.array([[annotation[4] * im_w, annotation[5] * im_h],
[annotation[6] * im_w, annotation[7] * im_h],
[annotation[8] * im_w, annotation[9] * im_h],
[annotation[10] * im_w, annotation[11] * im_h],
[annotation[12] * im_w, annotation[13] * im_h]],
dtype=np.float32)
# align the detected face
M, pose_index = estimate_norm(lm_align, 256, "arcface", shrink_factor=1.0)
im_aligned = (cv2.warpAffine(im, M, (256, 256), borderValue=0.0) - 127.5) / 127.5
if "adversarial defense" in settings:
eps = adv_slider / 200
X = tf.convert_to_tensor(np.expand_dims(im_aligned, axis=0))
with tf.GradientTape() as tape:
tape.watch(X)
X_z = ArcFaceE(tf.image.resize(X * 0.5 + 0.5, [112, 112]))
output = R([X, X_z])
loss = tf.reduce_mean(tf.abs(0 - output))
gradient = tf.sign(tape.gradient(loss, X))
adv_x = X + eps * gradient
im_aligned = tf.clip_by_value(adv_x, -1, 1)[0]
if "anonymize" in settings and "reconstruction attack" not in settings:
"""source_z = ArcFace.predict(np.expand_dims(tf.image.resize(im_aligned, [112, 112]) / 255.0, axis=0))
anon_ratio = int(512 * (slider / 100))
anon_vector = np.ones(shape=(1, 512))
anon_vector[:, :anon_ratio] = -1
np.random.shuffle(anon_vector)
source_z *= anon_vector"""
slider_weight = slider / 100
target_z = ArcFace.predict(np.expand_dims(tf.image.resize(im_aligned, [112, 112]) * 0.5 + 0.5, axis=0))
# source_z = IDP.predict(target_z)
source_z = slider_weight * source_z + (1 - slider_weight) * target_z
if "reconstruction attack" in settings:
source_z = ArcFaceE.predict(np.expand_dims(tf.image.resize(im_aligned, [112, 112]) * 0.5 + 0.5, axis=0))
# face swap
if "reconstruction attack" not in settings:
changed_face_cage = G.predict([np.expand_dims(im_aligned, axis=0),
source_z])
changed_face = changed_face_cage[0] * 0.5 + 0.5
# get inverse transformation landmarks
transformed_lmk = transform_landmark_points(M, lm_align)
# warp image back
iM, _ = inverse_estimate_norm(lm_align, transformed_lmk, 256, "arcface", shrink_factor=1.0)
iim_aligned = cv2.warpAffine(changed_face, iM, im_shape, borderValue=0.0)
# blend swapped face with target image
blend_mask = cv2.warpAffine(blend_mask_base, iM, im_shape, borderValue=0.0)
blend_mask = np.expand_dims(blend_mask, axis=-1)
total_img = (iim_aligned * blend_mask + total_img * (1 - blend_mask))
else:
changed_face_cage = R.predict([np.expand_dims(im_aligned, axis=0),
source_z])
changed_face = changed_face_cage[0] * 0.5 + 0.5
# get inverse transformation landmarks
transformed_lmk = transform_landmark_points(M, lm_align)
# warp image back
iM, _ = inverse_estimate_norm(lm_align, transformed_lmk, 256, "arcface", shrink_factor=1.0)
iim_aligned = cv2.warpAffine(changed_face, iM, im_shape, borderValue=0.0)
# blend swapped face with target image
blend_mask = cv2.warpAffine(blend_mask_base, iM, im_shape, borderValue=0.0)
blend_mask = np.expand_dims(blend_mask, axis=-1)
total_img = (iim_aligned * blend_mask + total_img * (1 - blend_mask))
if "compare" in settings:
total_img = np.concatenate((im / 255.0, total_img), axis=1)
total_img = np.clip(total_img, 0, 1)
total_img *= 255.0
total_img = total_img.astype('uint8')
return total_img
except Exception as e:
print(e)
return None
description = "Performs subject agnostic identity transfer from a source face to all target faces. \n\n" \
"Implementation and demo of FaceDancer, accepted to WACV 2023. \n\n" \
"Pre-print: https://arxiv.org/abs/2210.10473 \n\n" \
"Code: https://github.com/felixrosberg/FaceDancer \n\n" \
"\n\n" \
"Options:\n\n" \
"-Compare returns the target image concatenated with the results.\n\n" \
"-Anonymize will ignore the source image and perform an identity permutation of target faces.\n\n" \
"-Reconstruction attack will attempt to invert the face swap or the anonymization.\n\n" \
"-Adversarial defense will add a permutation noise that disrupts the reconstruction attack.\n\n" \
"NOTE: There is no guarantees with the anonymization process currently.\n\n" \
"NOTE: source image with too high resolution may not work properly!"
examples = [["assets/rick.jpg", "assets/musk.jpg", 100, 10, ["compare"]],
["assets/musk.jpg", "assets/musk.jpg", 100, 10, ["anonymize"]]]
article = """
Demo is based of recent research from my Ph.D work. Results expects to be published in the coming months.
"""
iface = gradio.Interface(run_inference,
[gradio.inputs.Image(shape=None, label='Target'),
gradio.inputs.Image(shape=None, label='Source'),
gradio.inputs.Slider(0, 100, default=100, label="Anonymization ratio (%)"),
gradio.inputs.Slider(0, 100, default=100, label="Adversarial defense ratio (%)"),
gradio.inputs.CheckboxGroup(["compare",
"anonymize",
"reconstruction attack",
"adversarial defense"],
label='Options')],
gradio.outputs.Image(),
title="Face Swap",
description=description,
examples=examples,
article=article,
layout="vertical")
iface.launch()