Akshita1414's picture
Create app.py
9f46ab9 verified
raw
history blame contribute delete
974 Bytes
import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch
model_name = "prithivMLmods/open-deepfake-detection"
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)
model.eval()
def predict(img):
inputs = processor(images=img, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probs = torch.softmax(logits, dim=1).squeeze()
is_ai = bool(torch.argmax(probs).item())
confidence = float(probs[1].item()) if is_ai else float(probs[0].item())
message = "AI-generated image detected." if is_ai else "Image appears original/authentic."
return {"is_ai_generated": is_ai, "confidence": confidence, "message": message}
iface = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs="json", allow_flagging="never")
iface.launch()