Update app.py

app.py

@@ -1,70 +1,254 @@
-# app.py (Option B - Minimal local pipeline; may use more RAM)
-import os, io, base64, traceback
 import gradio as gr
 from transformers import pipeline
-from PIL import Image
+from PIL import Image, ImageFilter, ImageOps
+import numpy as np
+import traceback
+import io
+import base64
 
-MODEL_ID = "Ateeqq/ai-vs-human-image-detector"
-HF_TOKEN = os.environ.get("HF_TOKEN")  # set if model private
+# -----------------------------
+# Your original model list
+# -----------------------------
+models = [
+    ("Ateeqq/ai-vs-human-image-detector", "ateeq"),
+    ("umm-maybe/AI-image-detector", "umm_maybe"),
+    ("dima806/ai_vs_human_generated_image_detection", "dimma"),
+]
 
-# Try to load pipeline (defensive)
+# load pipelines (same as your working code)
 pipes = []
-load_error = None
-try:
-    pipes.append((MODEL_ID, pipeline("image-classification", model=MODEL_ID, use_auth_token=HF_TOKEN)))
-    load_error = None
-    print(f"[INFO] Loaded {MODEL_ID}")
-except Exception as e:
-    load_error = repr(e)
-    print("[ERROR] Failed to load pipeline:", load_error)
-
-def predict(image: Image.Image):
-    if image is None:
-        return None, "<div style='color:red;'>Upload an image first</div>", load_error or ""
-    if not pipes:
-        # Show the exact load error to help debugging
-        return image, "<div style='color:red;'>No models loaded</div>", load_error or "No pipeline"
-    model_id, pipe = pipes[0]
+for model_id, _ in models:
     try:
-        res = pipe(image)
-        if not res:
-            return image, "<div style='color:red;'>Model returned no results</div>", ""
-        top = res[0]
-        label = top.get("label","").lower()
-        score = top.get("score", 0.0) * 100
+        pipes.append((model_id, pipeline("image-classification", model=model_id)))
+        print(f"Loaded {model_id}")
+    except Exception as e:
+        print(f"Error loading {model_id}: {e}")
+
+# -----------------------------
+# Helper: simple texture-based saliency map (no cv2, no model internals)
+# - This approximates "where the image has high-frequency detail"
+# - Not true Grad-CAM, but a lightweight explainability overlay that's safe to run in Spaces
+# -----------------------------
+def compute_texture_heatmap(pil_img, downsample=128):
+    """
+    Returns a 2D float numpy array (0..1) heatmap highlighting textured/high-frequency regions.
+    Steps:
+      - convert to grayscale
+      - blur to remove low-frequency shading
+      - compute absolute difference between original and blurred to highlight texture
+      - normalize
+    """
+    try:
+        # convert and resize for speed
+        w, h = pil_img.size
+        short = min(downsample, max(64, min(w, h)))
+        img_small = pil_img.convert("L").resize((short, short), resample=Image.BILINEAR)
+        # blurred version
+        blurred = img_small.filter(ImageFilter.GaussianBlur(radius=3))
+        # absolute difference
+        arr_orig = np.array(img_small).astype(np.float32) / 255.0
+        arr_blur = np.array(blurred).astype(np.float32) / 255.0
+        diff = np.abs(arr_orig - arr_blur)
+        # amplify small differences
+        diff = diff ** 0.8
+        # normalize to 0..1
+        diff = diff - diff.min()
+        diff = diff / (diff.max() + 1e-8)
+        return diff
+    except Exception as e:
+        print("compute_texture_heatmap error:", e)
+        return None
+
+def apply_colormap_numpy(heatmap):
+    """
+    Simple jet-like colormap without cv2.
+    heatmap: 2D float array 0..1
+    returns: HxWx3 uint8 RGB
+    """
+    h = np.clip(heatmap, 0.0, 1.0)
+    c = np.zeros((h.shape[0], h.shape[1], 3), dtype=np.float32)
+    c[..., 0] = np.clip(1.5 - 4.0 * np.abs(h - 0.25), 0, 1)  # R
+    c[..., 1] = np.clip(1.5 - 4.0 * np.abs(h - 0.5), 0, 1)   # G
+    c[..., 2] = np.clip(1.5 - 4.0 * np.abs(h - 0.75), 0, 1)  # B
+    return (c * 255).astype(np.uint8)
+
+def overlay_heatmap_on_pil(orig_pil, heatmap, alpha=0.55):
+    """
+    orig_pil: PIL RGB
+    heatmap: small 2D float array (0..1) -> will be resized to image
+    returns: PIL RGB overlay image
+    """
+    try:
+        orig = np.array(orig_pil.convert("RGB")).astype(np.uint8)
+        # resize heatmap to image size using PIL
+        hm_img = Image.fromarray((np.clip(heatmap,0,1) * 255).astype(np.uint8))
+        hm_resized = np.array(hm_img.resize((orig.shape[1], orig.shape[0]), resample=Image.BILINEAR)) / 255.0
+        colored = apply_colormap_numpy(hm_resized)
+        overlay = np.clip(orig * (1 - alpha) + colored * alpha, 0, 255).astype(np.uint8)
+        return Image.fromarray(overlay)
+    except Exception as e:
+        print("overlay_heatmap_on_pil error:", e)
+        return orig_pil
+
+# -----------------------------
+# Your original predict function, extended to return overlay + reason
+# -----------------------------
+def predict_image(image: Image.Image):
+    try:
+        results = []
+        for _, pipe in pipes:
+            # some pipelines may raise; make it robust
+            try:
+                res = pipe(image)
+                if isinstance(res, list) and res:
+                    res0 = res[0]
+                elif isinstance(res, dict):
+                    res0 = res
+                else:
+                    res0 = {"label":"error","score":0.0}
+            except Exception as e:
+                print("pipeline error:", e)
+                res0 = {"label":"error","score":0.0}
+            results.append(res0)
+
+        if not results:
+            return "<div style='color:red;'>No models loaded</div>", None, "no pipelines"
+
+        final_result = results[0]
+        label = final_result.get("label","").lower()
+        score = final_result.get("score",0.0) * 100
+
         if "ai" in label or "fake" in label:
             verdict = f"🧠 AI-Generated ({score:.1f}% confidence)"
             color = "#007BFF"
         else:
             verdict = f"🧍 Human-Made ({score:.1f}% confidence)"
             color = "#4CAF50"
+
+        # create the same styled HTML box you had
         html = f"""
-        <div style='background:linear-gradient(135deg,{color}33,#1a1a1a);
-                    border:2px solid {color}; border-radius:12px; padding:18px;
-                    text-align:center; color:white; font-weight:700;'>
-          {verdict}<div style="font-size:12px;opacity:0.85;margin-top:6px">Model: {model_id}</div>
+        <div class='result-box' style="
+            background: linear-gradient(135deg, {color}33, #1a1a1a);
+            border: 2px solid {color};
+            border-radius: 15px;
+            padding: 25px;
+            text-align: center;
+            color: white;
+            font-size: 20px;
+            font-weight: 600;
+            box-shadow: 0 0 20px {color}55;
+            animation: fadeIn 0.6s ease-in-out;
+        ">
+            {verdict}
         </div>
         """
-        return image, html, ""
+
+        # compute a lightweight texture heatmap (fast) and overlay
+        heatmap = compute_texture_heatmap(image, downsample=160)
+        overlay_img = None
+        explain_reason = ""
+        if heatmap is None:
+            explain_reason = "explainability failed"
+        else:
+            try:
+                overlay_img = overlay_heatmap_on_pil(image, heatmap, alpha=0.55)
+                explain_reason = "Texture-based saliency overlay (approximate explainability)"
+            except Exception as e:
+                print("overlay creation failed:", e)
+                overlay_img = None
+                explain_reason = "overlay failed"
+
+        # return: html string, overlay PIL image (or None), explain_reason text
+        return html, overlay_img, explain_reason
+
     except Exception as e:
-        err = repr(e)
-        return image, f"<div style='color:red;'>Inference failed: {err}</div>", err
+        traceback.print_exc()
+        return f"<div style='color:red;'>Error analyzing image: {str(e)}</div>", None, "error"
 
+# -----------------------------
+# CSS (same as yours)
+# -----------------------------
 css = """
-.gradio-container { font-family: 'Poppins', sans-serif; }
+body, .gradio-container {
+    font-family: 'Poppins', sans-serif !important;
+    background: transparent !important;
+}
+h1 {
+    text-align: center;
+    font-weight: 700;
+    color: #007BFF;
+    margin-bottom: 10px;
+}
+.gr-button-primary {
+    background-color: #007BFF !important;
+    color: white !important;
+    font-weight: 600;
+    border-radius: 10px;
+    height: 45px;
+}
+.gr-button-secondary {
+    background-color: #dc3545 !important;
+    color: white !important;
+    border-radius: 10px;
+    height: 45px;
+}
+#pulse-loader {
+    width: 100%;
+    height: 4px;
+    background: linear-gradient(90deg, #007BFF, #00C3FF);
+    animation: pulse 1.2s infinite ease-in-out;
+    border-radius: 2px;
+    box-shadow: 0 0 10px #007BFF;
+}
+@keyframes pulse {
+    0% { transform: scaleX(0.1); opacity: 0.6; }
+    50% { transform: scaleX(1); opacity: 1; }
+    100% { transform: scaleX(0.1); opacity: 0.6; }
+}
+@keyframes fadeIn {
+    from { opacity: 0; transform: scale(0.95); }
+    to { opacity: 1; transform: scale(1); }
+}
 """
 
-with gr.Blocks(css=css) as demo:
-    gr.Markdown("<h2>🔍 Unreal Eye (Local single-model)</h2>")
+# -----------------------------
+# Gradio UI (keeps your layout)
+# -----------------------------
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
+    gr.Markdown("<h1>🔍 AI Image Detector</h1>")
+
     with gr.Row():
-        with gr.Column():
-            inp = gr.Image(type="pil", label="Upload an image")
-            btn = gr.Button("Analyze")
-            btn_clear = gr.Button("Clear")
-        with gr.Column():
-            out_img = gr.Image(type="pil", label="Original / Overlay")
-            out_html = gr.HTML()
-            load_box = gr.Textbox(label="Load status / explainability", value=(load_error or "Model loaded" if pipes else "No model loaded"), interactive=False)
-    btn.click(predict, inputs=inp, outputs=[out_img, out_html, load_box])
-    btn_clear.click(lambda: (None, "", ""), outputs=[out_img, out_html, load_box])
+        with gr.Column(scale=1):
+            image_input = gr.Image(type="pil", label="Upload an image")
+            analyze_button = gr.Button("Analyze", variant="primary")
+            clear_button = gr.Button("Clear", variant="secondary")
+            loader = gr.HTML("")
+        with gr.Column(scale=1):
+            # show original / overlay side-by-side like you had
+            orig_display = gr.Image(type="pil", label="Upload an image")
+            overlay_display = gr.Image(type="pil", label="Original / Overlay")
+            explain_box = gr.Markdown("Explainability:")
+            explain_text = gr.Textbox(label="", interactive=False)
+
+            output = gr.HTML(label="Result")
+
+    def analyze(img):
+        if img is None:
+            return ("", None, None, "<div style='color:red;'>Please upload an image first!</div>")
+        loader_html = "<div id='pulse-loader'></div>"
+        yield (loader_html, None, None, "")  # show loader
+
+        # run prediction + explain
+        html, overlay_img, explain_reason = predict_image(img)
+
+        # if overlay exists, show both original and overlay
+        if overlay_img is not None:
+            yield ("", img, overlay_img, html + f"<div style='margin-top:8px; color:#ccc; font-size:12px;'>{explain_reason}</div>")
+        else:
+            # no overlay: show original and message
+            yield ("", img, img, html + f"<div style='margin-top:8px; color:#ccc; font-size:12px;'>{explain_reason}</div>")
+
+    analyze_button.click(analyze, inputs=image_input, outputs=[loader, orig_display, overlay_display, output])
+    clear_button.click(lambda: ("", None, None, ""), outputs=[loader, orig_display, overlay_display, output])
+
 demo.launch()