Spaces:
Sleeping
Sleeping
File size: 3,649 Bytes
391d6ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | import gradio as gr
import torch
import mimetypes
from PIL import Image
import cv2
from torchvision.models import efficientnet_b0
from torchvision import transforms
# =========================
# Load Model
# =========================
def load_model():
model = efficientnet_b0()
model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, 2)
model.load_state_dict(torch.load("models/best_model-v3.pt", map_location="cpu"))
model.eval()
return model
model = load_model()
# =========================
# Preprocessing
# =========================
preprocess = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
)
])
# =========================
# Image Prediction
# =========================
def predict_image(path):
img = Image.open(path).convert("RGB")
tensor = preprocess(img).unsqueeze(0)
with torch.no_grad():
out = model(tensor)
probs = torch.softmax(out, dim=1)[0]
conf, pred = torch.max(probs, dim=0)
label = "🟢 Real" if pred.item() == 0 else "🔴 Deepfake"
return label, f"{conf.item()*100:.2f}%", img
# =========================
# Video Prediction (Every 10th Frame)
# =========================
def predict_video(path):
cap = cv2.VideoCapture(path)
frame_count = 0
predictions = []
preview_img = None
while True:
ret, frame = cap.read()
if not ret:
break
# Process every 10th frame
if frame_count % 10 == 0:
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
img = Image.fromarray(frame_rgb)
# Save first sampled frame for preview
if preview_img is None:
preview_img = img
tensor = preprocess(img).unsqueeze(0)
with torch.no_grad():
out = model(tensor)
probs = torch.softmax(out, dim=1)[0]
predictions.append(probs)
frame_count += 1
cap.release()
if len(predictions) == 0:
return "❌ No valid frames found", "", None
# Average all frame probabilities
avg_probs = torch.stack(predictions).mean(dim=0)
conf, pred = torch.max(avg_probs, dim=0)
label = "🟢 Real (Multi-frame)" if pred.item() == 0 else "🔴 Deepfake (Multi-frame)"
return label, f"{conf.item()*100:.2f}%", preview_img
# =========================
# Main Prediction Router
# =========================
def predict_file(file_obj):
if file_obj is None:
return "⚠️ No file selected", "", None
path = file_obj.name
mime, _ = mimetypes.guess_type(path)
if mime and mime.startswith("image"):
return predict_image(path)
elif mime and mime.startswith("video"):
return predict_video(path)
else:
return "Unsupported file type", "", None
# =========================
# Gradio UI
# =========================
with gr.Blocks(title="Deepfake Detector") as demo:
gr.Markdown("## 🧠 Deepfake Detector\nUpload an image or video to analyze authenticity.")
file_input = gr.File(
label="Drop File Here",
file_types=[".jpg", ".jpeg", ".png", ".mp4", ".mov"]
)
with gr.Row():
prediction = gr.Textbox(label="Prediction", interactive=False)
confidence = gr.Textbox(label="Confidence (%)", interactive=False)
preview = gr.Image(label="Preview", interactive=False)
file_input.change(
fn=predict_file,
inputs=file_input,
outputs=[prediction, confidence, preview]
)
demo.launch() |