|
import spaces |
|
import gradio as gr |
|
import torch |
|
from PIL import Image |
|
import numpy as np |
|
from util.feature_extraction_utils import normalize_transforms |
|
from util.attack_utils import Attack |
|
from util.prepare_utils import prepare_models, prepare_dir_vec, get_ensemble |
|
from align.detector import detect_faces |
|
from align.align_trans import get_reference_facial_points, warp_and_crop_face |
|
import torchvision.transforms as transforms |
|
|
|
to_tensor = transforms.ToTensor() |
|
|
|
eps = 0.05 |
|
n_iters = 50 |
|
input_size = [112, 112] |
|
attack_type = "lpips" |
|
c_tv = None |
|
c_sim = 0.05 |
|
lr = 0.0025 |
|
net_type = "alex" |
|
noise_size = 0.005 |
|
n_starts = 1 |
|
kernel_size_gf = 7 |
|
sigma_gf = 3 |
|
combination = True |
|
using_subspace = False |
|
V_reduction_root = "./" |
|
model_backbones = ["IR_152", "IR_152", "ResNet_152", "ResNet_152"] |
|
model_roots = [ |
|
"https://github.com/cmu-spuds/lowkey_gradio/releases/download/weights/Backbone_IR_152_Arcface_Epoch_112.pth", |
|
"https://github.com/cmu-spuds/lowkey_gradio/releases/download/weights/Backbone_IR_152_Cosface_Epoch_70.pth", |
|
"https://github.com/cmu-spuds/lowkey_gradio/releases/download/weights/Backbone_ResNet_152_Arcface_Epoch_65.pth", |
|
"https://github.com/cmu-spuds/lowkey_gradio/releases/download/weights/Backbone_ResNet_152_Cosface_Epoch_68.pth", |
|
] |
|
direction = 1 |
|
crop_size = 112 |
|
scale = crop_size / 112.0 |
|
|
|
for root in model_roots: |
|
torch.hub.load_state_dict_from_url(root, map_location="cpu", progress=True) |
|
|
|
|
|
@spaces.GPU(duration=120) |
|
def protect(img, progress=gr.Progress(track_tqdm=True)): |
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
models_attack, V_reduction, dim = prepare_models( |
|
model_backbones, |
|
input_size, |
|
model_roots, |
|
kernel_size_gf, |
|
sigma_gf, |
|
combination, |
|
using_subspace, |
|
V_reduction_root, |
|
) |
|
|
|
img = Image.fromarray(img) |
|
reference = get_reference_facial_points(default_square=True) * scale |
|
h, w, c = np.array(img).shape |
|
|
|
_, landmarks = detect_faces(img) |
|
facial5points = [[landmarks[0][j], landmarks[0][j + 5]] for j in range(5)] |
|
|
|
_, tfm = warp_and_crop_face( |
|
np.array(img), facial5points, reference, crop_size=(crop_size, crop_size) |
|
) |
|
|
|
|
|
theta = normalize_transforms(tfm, w, h) |
|
tensor_img = to_tensor(img).unsqueeze(0).to(device) |
|
|
|
V_reduction = None |
|
dim = 512 |
|
|
|
|
|
dir_vec_extractor = get_ensemble( |
|
models=models_attack, |
|
sigma_gf=None, |
|
kernel_size_gf=None, |
|
combination=False, |
|
V_reduction=V_reduction, |
|
warp=True, |
|
theta_warp=theta, |
|
) |
|
dir_vec = prepare_dir_vec(dir_vec_extractor, tensor_img, dim, combination) |
|
|
|
img_attacked = tensor_img.clone() |
|
attack = Attack( |
|
models_attack, |
|
dim, |
|
attack_type, |
|
eps, |
|
c_sim, |
|
net_type, |
|
lr, |
|
n_iters, |
|
noise_size, |
|
n_starts, |
|
c_tv, |
|
sigma_gf, |
|
kernel_size_gf, |
|
combination, |
|
warp=True, |
|
theta_warp=theta, |
|
V_reduction=V_reduction, |
|
) |
|
img_attacked = attack.execute(tensor_img, dir_vec, direction).detach().cpu() |
|
|
|
img_attacked_pil = transforms.ToPILImage()(img_attacked[0]) |
|
return img_attacked_pil |
|
|
|
|
|
gr.Interface( |
|
fn=protect, |
|
inputs=gr.components.Image(height=512, width=512), |
|
outputs=gr.components.Image(type="pil"), |
|
allow_flagging="never", |
|
).launch(show_error=True, quiet=False, share=False) |
|
|
