Update app.py

app.py

@@ -1,353 +1,70 @@
-# app.py
+# app.py (Option B - Minimal local pipeline; may use more RAM)
+import os, io, base64, traceback
 import gradio as gr
-from transformers import pipeline, AutoImageProcessor, AutoModelForImageClassification
+from transformers import pipeline
 from PIL import Image
-import traceback, io, base64, time
-import torch
-import torch.nn.functional as F
-import numpy as np
 
-# -------- CONFIG --------
-HF_TOKEN = None  # set to "hf_xxx" if you use private models
-
-# List the models you want to use. Keep one model if you need lower memory.
-models = [
-    ("Ateeqq/ai-vs-human-image-detector", "ateeq"),
-    ("umm-maybe/AI-image-detector", "umm_maybe"),
-    ("dima806/ai_vs_human_generated_image_detection", "dimma"),
-]
-
-# ---------- Utilities: safe overlay WITHOUT OPENCV ----------
-def apply_colormap_numpy(heatmap):
-    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_no_cv(orig_pil, heatmap, alpha=0.45):
-    orig = np.array(orig_pil.convert("RGB"))
-    # resize heatmap to original size
-    heatmap_img = Image.fromarray((np.clip(heatmap,0,1)*255).astype(np.uint8))
-    heatmap_resized = np.array(heatmap_img.resize((orig.shape[1], orig.shape[0]), resample=Image.BILINEAR)) / 255.0
-    colored = apply_colormap_numpy(heatmap_resized)
-    overlay = (orig * (1 - alpha) + colored * alpha).astype(np.uint8)
-    return Image.fromarray(overlay)
-
-# ---------- Safe Grad-CAM helper (expects a helper object that returns heatmap) ----------
-def safe_gradcam_run(grad_helper, model, input_tensor):
-    """
-    grad_helper should be callable like: heatmap, class_idx, logits = grad_helper(input_tensor, class_idx)
-    This wrapper keeps things safe and returns tuple (heatmap_or_None, class_idx_or_None, logits_or_None, error_or_None)
-    """
-    try:
-        out = model(input_tensor.unsqueeze(0))
-        logits = out.logits if hasattr(out, "logits") else out
-        if logits is None:
-            return None, None, None, "no logits returned"
-        class_idx = int(torch.argmax(logits, dim=1).item())
-        # call grad helper (it may perform backward internally)
-        heatmap, idx, logits_tensor = grad_helper(input_tensor, class_idx)
-        if heatmap is None:
-            return None, class_idx, logits_tensor, "gradcam returned no heatmap"
-        return heatmap, class_idx, logits_tensor, None
-    except Exception as e:
-        return None, None, None, f"GradCAM error: {repr(e)}"
-
-# ---------- Safe ViT attention rollout ----------
-def safe_vit_attention_heatmap(processor, model, image: Image.Image):
-    try:
-        if processor is None or model is None:
-            return None, "processor or model missing"
-        inputs = processor(images=image, return_tensors="pt")
-        outputs = model(**inputs, output_attentions=True)
-        attentions = getattr(outputs, "attentions", None)
-        if not attentions:
-            return None, "no attentions in model output"
-        result = None
-        for attn in attentions:
-            a = attn[0].mean(0).detach().cpu().numpy()  # (seq, seq)
-            a = np.maximum(a, 0)
-            a = a / (a.sum(-1, keepdims=True) + 1e-8)
-            result = a if result is None else a @ result
-        cls_attn = result[0, 1:]
-        n_tokens = cls_attn.shape[0]
-        grid = int(np.round(np.sqrt(n_tokens)))
-        if grid * grid != n_tokens:
-            # best-effort reshape
-            grid = int(np.round(np.sqrt(n_tokens)))
-        heatmap = cls_attn.reshape(grid, grid)
-        heatmap = heatmap - heatmap.min()
-        heatmap = heatmap / (heatmap.max() + 1e-8)
-        return heatmap, None
-    except Exception as e:
-        return None, f"ViT rollout error: {repr(e)}"
-
-# ---------- Load pipelines and raw models (defensive) ----------
-pipes = []            # list of (model_id, pipeline)
-hf_models = {}        # model_id -> dict with processor/model/explain_type/helper
-
-for model_id, short in models:
-    # load inference pipeline (fast)
-    try:
-        p = pipeline("image-classification", model=model_id, use_auth_token=HF_TOKEN)
-        pipes.append((model_id, p))
-        print(f"[INFO] Loaded pipeline: {model_id}")
-    except Exception as e:
-        print(f"[WARN] Failed to load pipeline for {model_id}: {e}")
-
-    # try to load raw model + processor for explainability
-    proc = None
-    raw_model = None
-    explain_type = "none"
-    helper = None
-    try:
-        proc = AutoImageProcessor.from_pretrained(model_id, use_auth_token=HF_TOKEN)
-    except Exception:
-        # fallback: processor may not exist
-        proc = None
-
-    try:
-        raw_model = AutoModelForImageClassification.from_pretrained(model_id, use_auth_token=HF_TOKEN)
-        raw_model.eval()
-        # attempt to detect conv layers for Grad-CAM
-        base = None
-        for cand in ("base_model", "backbone", "model", "vit", "resnet", "conv_stem"):
-            if hasattr(raw_model, cand):
-                base = getattr(raw_model, cand)
-                break
-        if base is None:
-            base = raw_model
-
-        # find last conv2d if exists
-        last_conv = None
-        for name, m in base.named_modules():
-            if isinstance(m, torch.nn.Conv2d):
-                last_conv = m
-        if last_conv is not None:
-            explain_type = "gradcam"
-            # Create a small grad-cam helper object that registers hooks on last_conv
-            class GradCAMHelper:
-                def __init__(self, model, target_layer):
-                    self.model = model
-                    self.target_layer = target_layer
-                    self.activations = None
-                    self.gradients = None
-                    target_layer.register_forward_hook(self._save_activation)
-                    try:
-                        target_layer.register_backward_hook(self._save_gradient)
-                    except Exception:
-                        target_layer.register_full_backward_hook(self._save_gradient)
-
-                def _save_activation(self, module, inp, out):
-                    self.activations = out.detach()
-
-                def _save_gradient(self, module, grad_input, grad_output):
-                    self.gradients = grad_output[0].detach()
-
-                def __call__(self, input_tensor, class_idx=None):
-                    # forward
-                    out = self.model(input_tensor.unsqueeze(0))
-                    logits = out.logits if hasattr(out, "logits") else out
-                    if class_idx is None:
-                        class_idx = int(torch.argmax(logits, dim=1).item())
-                    self.model.zero_grad()
-                    score = logits[0, class_idx]
-                    score.backward(retain_graph=False)
-                    if self.gradients is None or self.activations is None:
-                        raise RuntimeError("gradcam hooks did not capture activations/gradients")
-                    pooled_grads = torch.mean(self.gradients[0], dim=(1,2))
-                    activ = self.activations[0].cpu()
-                    for i in range(activ.shape[0]):
-                        activ[i,:,:] *= pooled_grads[i].cpu()
-                    heatmap = torch.sum(activ, dim=0).cpu().numpy()
-                    heatmap = np.maximum(heatmap, 0)
-                    heatmap = heatmap - heatmap.min()
-                    heatmap = heatmap / (heatmap.max() + 1e-8)
-                    return heatmap, class_idx, logits
-            helper = GradCAMHelper(raw_model, last_conv)
-            print(f"[INFO] {model_id} -> gradcam ready")
-        else:
-            # if model looks like ViT (common in config.architectures)
-            cfg = getattr(raw_model, "config", None)
-            archs = getattr(cfg, "architectures", None) if cfg is not None else None
-            if archs and any("ViT" in a or "VisionTransformer" in a for a in archs):
-                explain_type = "vit"
-                helper = None
-                print(f"[INFO] {model_id} -> detected ViT, will use attention rollout")
-            else:
-                explain_type = "none"
-                helper = None
-                print(f"[INFO] {model_id} -> no explainability detected")
-    except Exception as e:
-        print(f"[WARN] Could not load raw HF model for explainability {model_id}: {e}")
-        raw_model = None
-        proc = proc
-        explain_type = "none"
-        helper = None
-
-    hf_models[model_id] = {
-        "processor": proc,
-        "model": raw_model,
-        "explain_type": explain_type,
-        "helper": helper
-    }
-
-# ---------- Prediction + explain wrapper ----------
-def predict_image_with_explain(image: Image.Image):
+MODEL_ID = "Ateeqq/ai-vs-human-image-detector"
+HF_TOKEN = os.environ.get("HF_TOKEN")  # set if model private
+
+# Try to load pipeline (defensive)
+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]
     try:
-        results = []
-        for model_id, pipe in pipes:
-            try:
-                res = pipe(image)
-                results.append((model_id, res[0] if isinstance(res, list) and res else {"label":"error","score":0.0}))
-            except Exception as e:
-                results.append((model_id, {"label":"error","score":0.0}))
-
-        if not results:
-            return {"html": "<div style='color:red;'>No models loaded</div>", "overlay": None, "explain_reason": "no pipelines"}
-
-        final_model_id, final_res = results[0]
-        label = final_res.get("label","").lower()
-        score = final_res.get("score",0.0) * 100
+        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
         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"
-
-        overlay_data_uri = None
-        explain_reason = ""
-
-        explain_entry = hf_models.get(final_model_id)
-        if explain_entry:
-            etype = explain_entry.get("explain_type","none")
-            try:
-                if etype == "gradcam" and explain_entry.get("helper") is not None:
-                    proc = explain_entry.get("processor")
-                    raw_model = explain_entry.get("model")
-                    # prepare input tensor robustly
-                    if proc is not None:
-                        inputs = proc(images=image, return_tensors="pt")
-                        # common key names
-                        input_tensor = inputs.get("pixel_values") or inputs.get("input_tensor") or list(inputs.values())[0]
-                        if isinstance(input_tensor, (list,tuple)):
-                            input_tensor = input_tensor[0]
-                        if isinstance(input_tensor, torch.Tensor) and input_tensor.dim()==4:
-                            input_tensor = input_tensor[0]
-                    else:
-                        # fallback preproc
-                        from torchvision import transforms
-                        pre = transforms.Compose([
-                            transforms.Resize((224,224)),
-                            transforms.ToTensor(),
-                            transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
-                        ])
-                        input_tensor = pre(image)
-                    heatmap, class_idx, logits, err = safe_gradcam_run(explain_entry["helper"], raw_model, input_tensor)
-                    if err:
-                        explain_reason = err
-                    elif heatmap is None:
-                        explain_reason = "gradcam returned no heatmap"
-                    else:
-                        overlay_img = overlay_heatmap_on_pil_no_cv(image, heatmap, alpha=0.45)
-                        buf = io.BytesIO()
-                        overlay_img.save(buf, format="PNG")
-                        overlay_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
-                        overlay_data_uri = "data:image/png;base64," + overlay_b64
-                        explain_reason = "Grad-CAM heatmap (activations)"
-                elif etype == "vit" and explain_entry.get("model") is not None:
-                    proc = explain_entry.get("processor")
-                    raw_model = explain_entry.get("model")
-                    heatmap, err = safe_vit_attention_heatmap(proc, raw_model, image)
-                    if err:
-                        explain_reason = err
-                    elif heatmap is None:
-                        explain_reason = "vit produced no heatmap"
-                    else:
-                        overlay_img = overlay_heatmap_on_pil_no_cv(image, heatmap, alpha=0.45)
-                        buf = io.BytesIO()
-                        overlay_img.save(buf, format="PNG")
-                        overlay_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
-                        overlay_data_uri = "data:image/png;base64," + overlay_b64
-                        explain_reason = "ViT attention rollout heatmap"
-                else:
-                    explain_reason = "No explainability available for this model"
-            except Exception as e:
-                explain_reason = f"Explain pipeline failed: {repr(e)}"
-        else:
-            explain_reason = "No raw HF entry for model"
-
         html = f"""
-        <div class='result-box' style="
-            background: linear-gradient(135deg, {color}33, #1a1a1a);
-            border: 2px solid {color};
-            border-radius: 15px;
-            padding: 20px;
-            text-align: center;
-            color: white;
-            font-size: 18px;
-            font-weight: 600;
-            box-shadow: 0 0 20px {color}55;
-            animation: fadeIn 0.6s ease-in-out;
-        ">
-            {verdict}
-            <div style="font-size:12px; margin-top:8px; font-weight:400; opacity:0.9;">
-                Model: <b>{final_model_id}</b> — Score: {score:.1f}%
-            </div>
+        <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>
         """
-
-        return {"html": html, "overlay": overlay_data_uri, "explain_reason": explain_reason}
+        return image, html, ""
     except Exception as e:
-        traceback.print_exc()
-        return {"html": f"<div style='color:red;'>Error analyzing image: {str(e)}</div>", "overlay": None, "explain_reason": ""}
+        err = repr(e)
+        return image, f"<div style='color:red;'>Inference failed: {err}</div>", err
 
-# ---------- Gradio UI ----------
 css = """
-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); } }
+.gradio-container { font-family: 'Poppins', sans-serif; }
 """
 
-with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
-    gr.Markdown("<h1>🔍 AI Image Detector w/ Explainability</h1>")
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("<h2>🔍 Unreal Eye (Local single-model)</h2>")
     with gr.Row():
-        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("")
-            gr.Markdown("Opacity:")
-            opacity = gr.Slider(minimum=0, maximum=1, value=0.6, step=0.05)
-        with gr.Column(scale=1):
-            image_display = gr.Image(type="pil", label="Original / Overlay", interactive=False)
-            output_html = gr.HTML(label="Result")
-            explanation_text = gr.Textbox(label="Explainability", interactive=False)
-
-    def analyze(img, op):
-        if img is None:
-            return (None, "<div style='color:red;'>Please upload an image first!</div>", "")
-        yield (None, "<div id='pulse-loader'></div>", "")
-        out = predict_image_with_explain(img)
-        overlay_uri = out.get("overlay")
-        if overlay_uri:
-            header, b64 = overlay_uri.split(",",1)
-            overlay_bytes = base64.b64decode(b64)
-            overlay_img = Image.open(io.BytesIO(overlay_bytes)).convert("RGB")
-        else:
-            overlay_img = img
-        explain_reason = out.get("explain_reason","")
-        html = out.get("html","")
-        yield (overlay_img, html, explain_reason)
-
-    analyze_button.click(analyze, inputs=[image_input, opacity], outputs=[image_display, output_html, explanation_text])
-    clear_button.click(lambda: (None, "", ""), outputs=[image_display, output_html, explanation_text])
-
+        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])
 demo.launch()