|
import os |
|
import sys |
|
from pathlib import Path |
|
import torch |
|
from PIL import Image |
|
import torchvision.transforms as transforms |
|
|
|
from networks.drn_seg import DRNSub, DRNSeg |
|
from utils.tools import * |
|
from utils.visualize import * |
|
|
|
import gradio as gr |
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
def load_classifier(type, model_path, device=torch.device("cpu")): |
|
if type == 'global': |
|
model = DRNSub(1) |
|
elif type == 'local': |
|
model = DRNSeg(2) |
|
state_dict = torch.load(model_path, map_location=device) |
|
model.load_state_dict(state_dict['model']) |
|
model.to(device) |
|
model.device = device |
|
model.eval() |
|
return model |
|
|
|
|
|
local_model_file = hf_hub_download( |
|
repo_id="radames/FALdetector", filename="local.pth", token=True) |
|
global_model_file = hf_hub_download( |
|
repo_id="radames/FALdetector", filename="global.pth", token=True) |
|
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
global_model = load_classifier("global", global_model_file, device) |
|
local_model = load_classifier("local", local_model_file, device) |
|
faces_model_file = 'utils/dlib_face_detector/mmod_human_face_detector.dat' |
|
|
|
tf = transforms.Compose([transforms.ToTensor(), |
|
transforms.Normalize(mean=[0.485, 0.456, 0.406], |
|
std=[0.229, 0.224, 0.225])]) |
|
|
|
|
|
def predict(img_path): |
|
im_w, im_h = Image.open(img_path).size |
|
faces = face_detection(img_path, verbose=False, |
|
model_file=faces_model_file) |
|
if len(faces) == 0: |
|
raise gr.Error("No face detected by dlib") |
|
|
|
face, box = faces[0] |
|
face = resize_shorter_side(face, 400)[0] |
|
face_tens = tf(face).to(device) |
|
|
|
with torch.no_grad(): |
|
prob = global_model(face_tens.unsqueeze(0))[0].sigmoid().cpu().item() |
|
flow = local_model(face_tens.unsqueeze(0))[0].cpu().numpy() |
|
flow = np.transpose(flow, (1, 2, 0)) |
|
h, w, _ = flow.shape |
|
|
|
|
|
modified = face.resize((w, h), Image.BICUBIC) |
|
modified_np = np.asarray(modified) |
|
reverse_np = warp(modified_np, flow) |
|
reverse = Image.fromarray(reverse_np) |
|
|
|
flow_magn = np.sqrt(flow[:, :, 0]**2 + flow[:, :, 1]**2) |
|
cv_out = get_heatmap_cv(modified_np, flow_magn, max_flow_mag=7) |
|
heat_map = Image.fromarray(cv_out) |
|
|
|
print(prob) |
|
return {"Probability FAL": prob}, modified, heat_map, reverse |
|
|
|
|
|
with gr.Blocks() as blocks: |
|
gr.Markdown(""" |
|
## Unofficial Demo |
|
### Detecting Photoshopped Faces by Scripting Photoshop |
|
#### FAL Detector Live Demo |
|
|
|
* https://arxiv.org/abs/1906.05856 |
|
* https://peterwang512.github.io/FALdetector/ |
|
""") |
|
with gr.Row(): |
|
with gr.Column(): |
|
in_image = gr.Image(label="Input Image", type="filepath") |
|
run_btn = gr.Button(label="Run") |
|
with gr.Column(): |
|
label = gr.Label( |
|
label="Probability being modified by Photoshop FAL") |
|
with gr.Row(): |
|
cropped = gr.Image(label="Cropped Input Image") |
|
heatmap = gr.Image(label="Heatmap") |
|
warped = gr.Image(label="Suggested Undo") |
|
run_btn.click(fn=predict, inputs=[in_image], outputs=[ |
|
label, cropped, heatmap, warped]) |
|
gr.Examples(fn=predict, |
|
examples=list(Path("./examples").glob("*.png")), |
|
inputs=[in_image], |
|
outputs=[label, cropped, heatmap, warped], |
|
cache_examples=True) |
|
blocks.launch() |
|
|
