Upload 4 files

app.py

@@ -0,0 +1,153 @@
+import gradio as gr
+
+# ---- IMPORT BACKENDS ----
+from web_backend import predict_image_pil
+from audio_inference import predict_audio
+
+
+# =========================
+# IMAGE LOGIC (UNCHANGED)
+# =========================
+def analyze_image(image):
+    label, confidence, heatmap = predict_image_pil(image)
+
+    if label == "Fake":
+        if confidence >= 90:
+            risk = "🚨 High likelihood of Deepfake"
+        elif confidence >= 60:
+            risk = "⚠️ Possibly Deepfake"
+        else:
+            risk = "⚠️ Uncertain Deepfake"
+    else:
+        if confidence >= 90:
+            risk = "✅ Likely Real"
+        elif confidence >= 60:
+            risk = "⚠️ Possibly Real"
+        else:
+            risk = "⚠️ Uncertain – Needs Review"
+
+    return label, f"{confidence} %", risk, heatmap
+
+
+# =========================
+# AUDIO LOGIC (UNCHANGED)
+# =========================
+def analyze_audio(audio_path):
+    label, confidence = predict_audio(audio_path)
+
+    if label == "fake":
+        if confidence >= 90:
+            risk = "🚨 High likelihood of Deepfake"
+        elif confidence >= 60:
+            risk = "⚠️ Possibly Deepfake"
+        else:
+            risk = "⚠️ Uncertain – Needs Review"
+    else:
+        if confidence >= 90:
+            risk = "✅ Likely Real"
+        elif confidence >= 60:
+            risk = "⚠️ Possibly Real"
+        else:
+            risk = "⚠️ Uncertain – Needs Review"
+
+    return label.capitalize(), f"{confidence} %", risk
+
+
+# =========================
+# UI
+# =========================
+with gr.Blocks() as demo:
+    gr.Markdown("# 🧠 Unified Deepfake Detection System")
+
+    with gr.Tabs():
+
+        # =====================
+        # HOME TAB
+        # =====================
+        with gr.Tab("🏠 Home"):
+            gr.Markdown(
+                """
+                ## Welcome 👋  
+                Select the type of media you want to analyze:
+                """
+            )
+
+            gr.Markdown("### 🔍 Choose Detection Mode")
+            gr.Markdown("- 🖼 **Image Deepfake Detection**\n- 🎧 **Audio Deepfake Detection**")
+
+            gr.Markdown(
+                """
+                👉 Use the tabs above to switch between Image and Audio detection.
+                """
+            )
+
+        # =====================
+        # IMAGE TAB
+        # =====================
+        with gr.Tab("🖼 Image Deepfake"):
+            gr.Markdown("# 🖼 Deepfake Image Detection System")
+
+            with gr.Row():
+                with gr.Column(scale=1):
+                    image_input = gr.Image(
+                        label="Upload Image",
+                        type="pil",
+                        height=280
+                    )
+                    img_submit = gr.Button("Submit")
+                    img_clear = gr.Button("Clear")
+
+                with gr.Column(scale=2):
+                    img_pred = gr.Text(label="Prediction")
+                    img_conf = gr.Text(label="Confidence")
+                    img_risk = gr.Text(label="Risk Assessment")
+                    img_heatmap = gr.Image(
+                        label="Explainability Heatmap",
+                        height=280
+                    )
+
+            img_submit.click(
+                fn=analyze_image,
+                inputs=image_input,
+                outputs=[img_pred, img_conf, img_risk, img_heatmap]
+            )
+
+            img_clear.click(
+                fn=lambda: (None, "", "", None),
+                inputs=None,
+                outputs=[image_input, img_pred, img_conf, img_risk]
+            )
+
+        # =====================
+        # AUDIO TAB
+        # =====================
+        with gr.Tab("🎧 Audio Deepfake"):
+            gr.Markdown("# 🎧 Deepfake Audio Detection System")
+
+            with gr.Row():
+                with gr.Column(scale=1):
+                    audio_input = gr.Audio(
+                        label="Upload Audio (.wav)",
+                        type="filepath"
+                    )
+                    aud_submit = gr.Button("Submit")
+                    aud_clear = gr.Button("Clear")
+
+                with gr.Column(scale=2):
+                    aud_pred = gr.Text(label="Prediction")
+                    aud_conf = gr.Text(label="Confidence")
+                    aud_risk = gr.Text(label="Risk Assessment")
+
+            aud_submit.click(
+                fn=analyze_audio,
+                inputs=audio_input,
+                outputs=[aud_pred, aud_conf, aud_risk]
+            )
+
+            aud_clear.click(
+                fn=lambda: (None, "", ""),
+                inputs=None,
+                outputs=[audio_input, aud_pred, aud_conf]
+            )
+
+demo.launch()

audio_inference.py

@@ -0,0 +1,57 @@
+import torch
+import librosa
+import numpy as np
+from transformers import Wav2Vec2Processor, Wav2Vec2ForSequenceClassification
+
+# =====================
+# CONFIG
+# =====================
+MODEL_DIR = "exported_audio_model"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+SR = 16000
+MAX_SAMPLES = 8 * SR  # 8 seconds
+
+# =====================
+# LOAD MODEL + PROCESSOR (ONCE)
+# =====================
+processor = Wav2Vec2Processor.from_pretrained(MODEL_DIR)
+model = Wav2Vec2ForSequenceClassification.from_pretrained(MODEL_DIR)
+model.to(DEVICE)
+model.eval()
+
+# =====================
+# PREDICT FUNCTION
+# =====================
+def predict_audio(wav_path):
+    # Load audio
+    audio, sr = librosa.load(wav_path, sr=SR, mono=True)
+
+    # Truncate if needed
+    if len(audio) > MAX_SAMPLES:
+        audio = audio[:MAX_SAMPLES]
+
+    # Processor handles padding
+    inputs = processor(
+        audio,
+        sampling_rate=SR,
+        return_tensors="pt",
+        padding=True,
+        return_attention_mask=True
+    )
+
+    input_values = inputs.input_values.to(DEVICE)
+    attention_mask = inputs.attention_mask.to(DEVICE)
+
+    with torch.no_grad():
+        outputs = model(
+            input_values=input_values,
+            attention_mask=attention_mask
+        )
+
+        probs = torch.softmax(outputs.logits, dim=1)[0]
+        pred_id = torch.argmax(probs).item()
+
+    label = model.config.id2label[pred_id]
+    confidence = probs[pred_id].item() * 100
+
+    return label, round(confidence, 2)

web_backend.py

@@ -0,0 +1,101 @@
+import torch
+from torchvision import transforms
+from transformers import ViTForImageClassification, ViTConfig
+from PIL import Image
+import numpy as np
+import matplotlib.pyplot as plt
+import io
+import os
+
+# -----------------------------
+# Device
+# -----------------------------
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# -----------------------------
+# Model Setup (SAME AS CMD)
+# -----------------------------
+config = ViTConfig.from_pretrained(
+    "google/vit-base-patch16-224",
+    num_labels=2,
+    output_attentions=True
+)
+
+model = ViTForImageClassification.from_pretrained(
+    "google/vit-base-patch16-224",
+    config=config,
+    ignore_mismatched_sizes=True
+)
+
+if os.path.exists("model/vit_real_fake_best.pth"):
+    model.load_state_dict(
+        torch.load("model/vit_real_fake_best.pth", map_location=device)
+    )
+
+model.to(device)
+model.eval()
+
+# -----------------------------
+# Image Preprocessing (IDENTICAL)
+# -----------------------------
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(
+        [0.485, 0.456, 0.406],
+        [0.229, 0.224, 0.225]
+    )
+])
+
+# -----------------------------
+# Attention Heatmap (IDENTICAL)
+# -----------------------------
+def get_attention_map(model, img_tensor):
+    with torch.no_grad():
+        outputs = model(img_tensor, output_attentions=True)
+        attn = outputs.attentions[-1].mean(dim=1)[0]
+        cls_attn = attn[0, 1:]
+
+        grid_size = int(cls_attn.size(0) ** 0.5)
+        cls_attn = cls_attn.reshape(grid_size, grid_size).cpu().numpy()
+        cls_attn = (cls_attn - cls_attn.min()) / (cls_attn.max() - cls_attn.min())
+
+        return cls_attn
+
+def overlay_heatmap_on_image(image, heatmap):
+    heatmap = np.uint8(255 * heatmap)
+    heatmap = Image.fromarray(heatmap).resize(image.size)
+    heatmap_np = np.array(heatmap)
+
+    fig, ax = plt.subplots(figsize=(4, 4))
+    ax.imshow(image)
+    ax.imshow(heatmap_np, cmap="jet", alpha=0.5)
+    ax.axis("off")
+
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
+    plt.close(fig)
+    buf.seek(0)
+
+    return Image.open(buf)
+
+# -----------------------------
+# Prediction Function (SOURCE OF TRUTH)
+# -----------------------------
+def predict_image_pil(image):
+    image = image.convert("RGB")
+    input_tensor = transform(image).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        outputs = model(input_tensor)
+        logits = outputs.logits
+        pred = torch.argmax(logits, dim=1).item()
+
+    label = "Fake" if pred == 0 else "Real"
+
+    attn_map = get_attention_map(model, input_tensor)
+    heatmap_img = overlay_heatmap_on_image(image, attn_map)
+
+    confidence = torch.softmax(logits, dim=1)[0][pred].item() * 100
+
+    return label, round(confidence, 2), heatmap_img