brendona17's picture
Update app.py
c3dc5dd verified
import gradio as gr
from transformers import AutoImageProcessor, SiglipForImageClassification
from PIL import Image
import torch
# ============================================================
# LOAD MODEL FROM HUGGING FACE HUB
# ============================================================
model_name = "prithivMLmods/Deepfake-Detect-Siglip2"
model = SiglipForImageClassification.from_pretrained(model_name)
processor = AutoImageProcessor.from_pretrained(model_name)
model.eval()
# ============================================================
# PREDICTION
# ============================================================
def predict(image):
if image is None:
return {"Error": "Upload an image"}
if not isinstance(image, Image.Image):
image = Image.fromarray(image)
image = image.convert("RGB")
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
probs = torch.nn.functional.softmax(outputs.logits, dim=1).squeeze().tolist()
labels = model.config.id2label
result = {labels[i]: round(probs[i], 4) for i in range(len(probs))}
print(f"Result: {result}")
return result
# ============================================================
# INTERFACE
# ============================================================
demo = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=2),
title="Deepfake Detection",
description="Upload an image to detect if it's Real or Fake"
)
if __name__ == "__main__":
demo.launch()