not-lain's picture
switched to pipeline
51abdcc
raw
history blame
1.47 kB
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
from transformers import pipeline
with zipfile.ZipFile("examples.zip","r") as zip_ref:
zip_ref.extractall(".")
pipe = pipeline(model="not-lain/deepfake",trust_remote_code=True)
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) # shuffle
def predict(input_image:Image.Image, true_label:str):
out = pipe.predict(input_image)
confidences,face_with_mask = out["confidences"], out["face_with_mask"]
return confidences, true_label, face_with_mask
interface = gr.Interface(
fn=predict,
inputs=[
gr.Image(label="Input Image", type="filepath"),
"text"
],
outputs=[
gr.Label(label="Class"),
"text",
gr.Image(label="Face with Explainability")
],
examples=[[examples[i]["path"], examples[i]["label"]] for i in range(10)]
).launch()