Update app.py

app.py

@@ -6,9 +6,10 @@ import numpy as np
 import cv2
 from PIL import Image
 import requests
-from torchvision import transforms
 import io
 import os
+from torchvision import transforms
+from transformers import pipeline
 
 # ── Model Definition ────────────────────────────────────────
 class CIFAKECNN(nn.Module):
@@ -32,12 +33,15 @@ class CIFAKECNN(nn.Module):
         x = torch.sigmoid(self.fc2(x))
         return x
 
-# ── Load Model ───────────────────────────────────────────────
+# ── Load our CNN ─────────────────────────────────────────────
 device = torch.device('cpu')
 model = CIFAKECNN().to(device)
 model.load_state_dict(torch.load('cnn_model.pth', map_location=device))
 model.eval()
 
+# ── Load pretrained detector ─────────────────────────────────
+pretrained_detector = pipeline("image-classification", model="Organika/sdxl-detector")
+
 # ── Transform ────────────────────────────────────────────────
 transform = transforms.Compose([
     transforms.Resize((32, 32)),
@@ -88,61 +92,74 @@ def get_gradcam_overlay(image_pil, image_tensor):
 # ── Detection Function ───────────────────────────────────────
 def detect_image(image):
     if image is None:
-        return "Please upload an image.", None
+        return "Please upload an image.", "Please upload an image.", None
     image_pil = Image.fromarray(image).convert('RGB')
     image_tensor = transform(image_pil)
+
+    # Our CNN
     with torch.no_grad():
         output = model(image_tensor.unsqueeze(0))
         confidence = output.item()
     label = "FAKE (AI-Generated)" if confidence > 0.5 else "REAL"
     confidence_pct = confidence if confidence > 0.5 else 1 - confidence
-    result = f"**{label}**\nConfidence: {confidence_pct:.1%}"
+    our_result = f"**{label}**\nConfidence: {confidence_pct:.1%}"
+
+    # Pretrained detector
+    pretrained_result = pretrained_detector(image_pil)
+    top = pretrained_result[0]
+    pretrained_label = top['label']
+    pretrained_conf = top['score']
+    pretrained_out = f"**{pretrained_label}**\nConfidence: {pretrained_conf:.1%}"
+
     gradcam_img = get_gradcam_overlay(image_pil, image_tensor)
-    return result, gradcam_img
+    return our_result, pretrained_out, gradcam_img
 
 # ── Generate & Detect Function ───────────────────────────────
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 def generate_and_detect(prompt):
     if not prompt:
-        return None, "Please enter a prompt.", None
+        return None, "Please enter a prompt.", "Please enter a prompt.", None
     response = requests.post(
         "https://router.huggingface.co/hf-inference/models/black-forest-labs/FLUX.1-schnell",
         headers={"Authorization": f"Bearer {HF_TOKEN}"},
         json={"inputs": prompt}
     )
     if response.status_code != 200:
-        return None, f"Generation failed: {response.text}", None
+        return None, f"Generation failed: {response.text}", "", None
     image_pil = Image.open(io.BytesIO(response.content)).convert('RGB')
-    result, gradcam_img = detect_image(np.array(image_pil))
-    return image_pil, result, gradcam_img
+    our_result, pretrained_out, gradcam_img = detect_image(np.array(image_pil))
+    return image_pil, our_result, pretrained_out, gradcam_img
 
 # ── Gradio Interface ─────────────────────────────────────────
 with gr.Blocks(title="AI Image Detector") as demo:
     gr.Markdown("# 🔍 Truth in the Noise: AI Image Detector")
-    gr.Markdown("Detect whether an image is real or AI-generated using our CNN model trained on CIFAKE.")
+    gr.Markdown("Detect whether an image is real or AI-generated. We compare our custom CNN (trained on CIFAKE) with a pretrained detector.")
 
     with gr.Tabs():
         with gr.Tab("📤 Upload & Detect"):
-            gr.Markdown("Upload any image to check if it's real or AI-generated.")
+            gr.Markdown("Upload any image to compare both models.")
             with gr.Row():
                 upload_input = gr.Image(label="Upload Image")
                 with gr.Column():
-                    upload_result = gr.Markdown(label="Result")
-                    upload_gradcam = gr.Image(label="GradCAM Visualization")
+                    upload_our = gr.Markdown(label="Our CNN (CIFAKE)")
+                    upload_pretrained = gr.Markdown(label="Pretrained Detector")
+                    upload_gradcam = gr.Image(label="GradCAM (Our CNN)")
             upload_btn = gr.Button("Detect", variant="primary")
-            upload_btn.click(detect_image, inputs=upload_input, outputs=[upload_result, upload_gradcam])
+            upload_btn.click(detect_image, inputs=upload_input,
+                           outputs=[upload_our, upload_pretrained, upload_gradcam])
 
         with gr.Tab("🎨 Generate & Detect"):
-            gr.Markdown("Enter a prompt to generate an AI image, then automatically detect it.")
+            gr.Markdown("Generate an AI image and detect it with both models.")
             prompt_input = gr.Textbox(label="Prompt", placeholder="e.g. a cat sitting on a chair")
             generate_btn = gr.Button("Generate & Detect", variant="primary")
             with gr.Row():
                 generated_img = gr.Image(label="Generated Image")
                 with gr.Column():
-                    generate_result = gr.Markdown(label="Result")
-                    generate_gradcam = gr.Image(label="GradCAM Visualization")
+                    generate_our = gr.Markdown(label="Our CNN (CIFAKE)")
+                    generate_pretrained = gr.Markdown(label="Pretrained Detector")
+                    generate_gradcam = gr.Image(label="GradCAM (Our CNN)")
             generate_btn.click(generate_and_detect, inputs=prompt_input,
-                             outputs=[generated_img, generate_result, generate_gradcam])
+                             outputs=[generated_img, generate_our, generate_pretrained, generate_gradcam])
 
 demo.launch()