|
import gradio as gr |
|
import torch |
|
import torch.nn.functional as F |
|
from facenet_pytorch import MTCNN, InceptionResnetV1 |
|
import os |
|
import numpy as np |
|
from PIL import Image |
|
import zipfile |
|
import cv2 |
|
from pytorch_grad_cam import GradCAM |
|
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget |
|
from pytorch_grad_cam.utils.image import show_cam_on_image |
|
|
|
with zipfile.ZipFile("examples.zip","r") as zip_ref: |
|
zip_ref.extractall(".") |
|
|
|
DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu' |
|
|
|
mtcnn = MTCNN( |
|
select_largest=False, |
|
post_process=False, |
|
device=DEVICE |
|
).to(DEVICE).eval() |
|
|
|
model = InceptionResnetV1( |
|
pretrained="vggface2", |
|
classify=True, |
|
num_classes=1, |
|
device=DEVICE |
|
) |
|
|
|
checkpoint = torch.load("resnetinceptionv1_epoch_32.pth", map_location=torch.device('cpu')) |
|
model.load_state_dict(checkpoint['model_state_dict']) |
|
model.to(DEVICE) |
|
model.eval() |
|
|
|
EXAMPLES_FOLDER = 'examples' |
|
examples_names = os.listdir(EXAMPLES_FOLDER) |
|
examples = [] |
|
for example_name in examples_names: |
|
example_path = os.path.join(EXAMPLES_FOLDER, example_name) |
|
label = example_name.split('_')[0] |
|
example = { |
|
'path': example_path, |
|
'label': label |
|
} |
|
examples.append(example) |
|
np.random.shuffle(examples) |
|
|
|
def predict(input_image:Image.Image, true_label:str): |
|
"""Predict the label of the input_image""" |
|
face = mtcnn(input_image) |
|
if face is None: |
|
raise Exception('No face detected') |
|
face = face.unsqueeze(0) |
|
face = F.interpolate(face, size=(256, 256), mode='bilinear', align_corners=False) |
|
|
|
|
|
prev_face = face.squeeze(0).permute(1, 2, 0).cpu().detach().int().numpy() |
|
prev_face = prev_face.astype('uint8') |
|
|
|
face = face.to(DEVICE) |
|
face = face.to(torch.float32) |
|
face = face / 255.0 |
|
face_image_to_plot = face.squeeze(0).permute(1, 2, 0).cpu().detach().int().numpy() |
|
|
|
target_layers=[model.block8.branch1[-1]] |
|
cam = GradCAM(model=model, target_layers=target_layers) |
|
targets = [ClassifierOutputTarget(0)] |
|
|
|
grayscale_cam = cam(input_tensor=face, targets=targets, eigen_smooth=True) |
|
grayscale_cam = grayscale_cam[0, :] |
|
visualization = show_cam_on_image(face_image_to_plot, grayscale_cam, use_rgb=True) |
|
face_with_mask = cv2.addWeighted(prev_face, 1, visualization, 0.5, 0) |
|
|
|
with torch.no_grad(): |
|
output = torch.sigmoid(model(face).squeeze(0)) |
|
prediction = "real" if output.item() < 0.5 else "fake" |
|
|
|
real_prediction = 1 - output.item() |
|
fake_prediction = output.item() |
|
|
|
confidences = { |
|
'real': real_prediction, |
|
'fake': fake_prediction |
|
} |
|
return confidences, true_label, face_with_mask |
|
|
|
interface = gr.Interface( |
|
fn=predict, |
|
inputs=[ |
|
gr.components.Image(label="Input Image", type="pil"), |
|
gr.components.Text(label="Your Text Input") |
|
], |
|
outputs=[ |
|
gr.components.Label(label="Class"), |
|
gr.components.Text(label="Your Text Output"), |
|
gr.components.Image(label="Face with Explainability", type="numpy") |
|
], |
|
examples=[[examples[i]["path"], examples[i]["label"]] for i in range(10)], |
|
cache_examples=True |
|
).launch() |
