Shiwanni's picture
Update app.py
e72bd1c verified
raw
history blame
1.06 kB
from transformers import ViTForImageClassification, ViTImageProcessor
import torch
from PIL import Image
import gradio as gr
# Load pre-trained model and processor
model_name = "facebook/deit-base-distilled-patch16-224"
processor = ViTImageProcessor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)
def detect_deepfake(image):
# Preprocess the image
inputs = processor(images=image, return_tensors="pt")
# Make prediction
outputs = model(**inputs)
logits = outputs.logits
predicted_class_idx = logits.argmax(-1).item()
# For demonstration, we'll assume class 0 is real and 1 is fake
# (In a real project, you'd need to verify this with your model)
return "Real" if predicted_class_idx == 0 else "Fake (Possible Deepfake)"
# Create a simple interface
iface = gr.Interface(
fn=detect_deepfake,
inputs=gr.Image(type="pil"),
outputs="text",
title="Deepfake Detection",
description="Upload an image to check if it might be a deepfake."
)
iface.launch()