Update app.py

app.py

@@ -1,145 +1,112 @@
 # app.py
+
+import gradio as gr
+from huggingface_hub import InferenceClient
+from PIL import Image
 import io
 import os
 import numpy as np
-from PIL import Image
-import gradio as gr
-from huggingface_hub import InferenceClient
 
-# Load API token from HF secrets
-HF_API_TOKEN = os.environ.get("HF_API_TOKEN", None)
+# -----------------------------
+# Hugging Face API Setup
+# -----------------------------
+HF_API_TOKEN = os.environ.get("HF_API_TOKEN")  # Add your token in Hugging Face Secrets
 client = InferenceClient(token=HF_API_TOKEN)
 
-# Models
-CLASSIFIER_MODEL = "prithivMLmods/deepfake-detector-model-v1"
-FORGERY_MODEL = "zhipeixu/fakeshield-v1-22b"
-
-def run_classification(img):
-    """ Deepfake / AI image detection model """
+# -----------------------------
+# Model Names
+# -----------------------------
+MODEL_1 = "prithivMLmods/deepfake-detector-model-v1"   # Deepfake detector
+MODEL_2 = "microsoft/dit-base-finetuned-aigc-detection"  # AIGC detector
+MODEL_3 = "zhipeixu/fakeshield-v1-22b"                 # Forgery detector
+
+# -----------------------------
+# Helper function: overlay mask on image
+# -----------------------------
+def overlay_mask(image, mask):
+    if mask is None:
+        return image
+    mask = np.array(mask.convert("L"))  # Convert mask to grayscale
+    mask = (mask > 128).astype(np.uint8) * 255  # Binary mask
+    overlay = Image.new("RGBA", image.size, (255,0,0,100))  # Red overlay
+    img_rgba = image.convert("RGBA")
+    img_rgba.paste(overlay, mask=Image.fromarray(mask))
+    return img_rgba
+
+# -----------------------------
+# Main function: Analyze image using 3 models
+# -----------------------------
+def analyze_image(image):
+    buf = io.BytesIO()
+    image.save(buf, format="PNG")
+    buf.seek(0)
+
+    # -------- MODEL 1: Deepfake Detector --------
     try:
-        buf = io.BytesIO()
-        img.save(buf, format="PNG")
-        buf.seek(0)
-
-        out = client.image_classification(
-            model=CLASSIFIER_MODEL,
-            inputs=buf
-        )
-
-        if isinstance(out, list) and len(out) > 0:
-            top = out[0]
-            label = top.get("label", "")
-            score = float(top.get("score", 0.0))
-
-            if "fake" in label.lower() or "ai" in label.lower():
-                return "AI-Generated", round(score * 100, 2)
-            else:
-                return "Real Image", round(score * 100, 2)
-
-        return "Unknown", 0.0
-
-    except Exception as e:
-        return f"Error: {e}", 0.0
+        out1 = client.image_classification(model=MODEL_1, inputs=buf)
+        label1 = out1[0]["label"]
+        score1 = round(out1[0]["score"] * 100, 2)
+    except:
+        label1, score1 = "Error", 0
+    buf.seek(0)
 
-def run_forgery_model(img):
-    """ Forgery & manipulation detection model """
+    # -------- MODEL 2: AIGC Detector --------
     try:
-        buf = io.BytesIO()
-        img.save(buf, format="PNG")
-        buf.seek(0)
-
-        out = client(
-            model=FORGERY_MODEL,
-            inputs=buf
-        )
-
-        result = {
-            "explanation": None,
-            "mask": None,
-            "raw": out
-        }
-
-        # Modern HF models return dict
-        if isinstance(out, dict):
-            result["explanation"] = out.get("explanation") or out.get("text")
-            result["mask"] = out.get("mask")
-            return result
-
-        # If output is a list of tokens
-        if isinstance(out, list):
-            explanation = []
-            for item in out:
-                if isinstance(item, dict):
-                    explanation.append(item.get("text") or item.get("label", ""))
-                elif isinstance(item, str):
-                    explanation.append(item)
-            result["explanation"] = " ".join(explanation)
-            return result
-
-        return result
-
-    except Exception as e:
-        return {"explanation": f"Error: {e}", "mask": None}
-
-def overlay_mask(img, mask_data):
-    """ Creates a red overlay on manipulated regions """
-    if mask_data is None:
-        return None
+        out2 = client.image_classification(model=MODEL_2, inputs=buf)
+        label2 = out2[0]["label"]
+        score2 = round(out2[0]["score"] * 100, 2)
+    except:
+        label2, score2 = "Error", 0
+    buf.seek(0)
 
+    # -------- MODEL 3: Forgery / Mask Detector --------
     try:
-        arr = np.array(mask_data)
-        if arr.max() <= 1:
-            arr = (arr * 255).astype("uint8")
-
-        mask = Image.fromarray(arr).resize(img.size).convert("L")
-        red = Image.new("RGBA", img.size, (255, 0, 0, 120))
-        overlay = Image.composite(red, Image.new("RGBA", img.size), mask)
-        final = Image.alpha_composite(img.convert("RGBA"), overlay)
-        return final
-
+        out3 = client(inputs=buf, model=MODEL_3)
+        explanation = out3.get("explanation", "No manipulation detected")
+        mask = out3.get("mask", None)
     except:
-        return None
-
-def analyze(image):
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image)
-
-    label, percent = run_classification(image)
-    forg = run_forgery_model(image)
-
-    explanation = forg.get("explanation") or "No clear manipulation detected."
-    mask = forg.get("mask")
-    overlay_image = overlay_mask(image, mask)
-
-    return image, f"{label} ({percent}%)", explanation, overlay_image
-
-
-# ---------------- UI --------------------
-title = "AI DeepFake & Manipulation Detector"
-description = """
-Upload an image to detect if it's AI-generated or manipulated.
-Two AI models are used:
-- Deepfake classifier (Real vs AI)
-- Forgery detector (Manipulated region + explanation)
-"""
+        explanation, mask = "Error detecting forgery", None
+
+    # -------- FINAL DECISION (Fusion) --------
+    ai_votes = 0
+    if "fake" in label1.lower() or "ai" in label1.lower():
+        ai_votes += 1
+    if "ai" in label2.lower() or "generated" in label2.lower():
+        ai_votes += 1
+
+    if ai_votes == 2:
+        final_label = "AI-GENERATED"
+    elif ai_votes == 1:
+        final_label = "Possibly AI-GENERATED"
+    else:
+        final_label = "REAL IMAGE"
+
+    # Overlay mask if exists
+    output_image = overlay_mask(image, mask)
+
+    # Return outputs
+    return (
+        output_image,
+        f"{final_label}",
+        f"Deepfake Model: {label1} ({score1}%)\nAIGC Model: {label2} ({score2}%)\nForgery Detector: {explanation}"
+    )
 
+# -----------------------------
+# Gradio Interface
+# -----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown(f"# {title}\n{description}")
-
+    gr.Markdown("<h2 style='text-align:center'>AI DeepFake & Manipulation Detector</h2>")
+    
     with gr.Row():
         inp = gr.Image(type="pil", label="Upload Image")
-        with gr.Column():
-            original_out = gr.Image(label="Original Image")
-            overlay_out = gr.Image(label="Manipulation Overlay")
-
-    label_out = gr.Textbox(label="Classification", interactive=False)
-    explanation_out = gr.Textbox(label="Manipulation Explanation", interactive=False)
-
-    btn = gr.Button("Analyze")
-    btn.click(
-        fn=analyze,
-        inputs=[inp],
-        outputs=[original_out, label_out, explanation_out, overlay_out]
-    )
-
+        out_img = gr.Image(type="pil", label="Result Image with Mask Overlay")
+    
+    out_text = gr.Textbox(label="Detection Result & Explanation", lines=8)
+    
+    btn = gr.Button("Analyze Image")
+    btn.click(fn=analyze_image, inputs=[inp], outputs=[out_img, out_text, out_text])
+
+# -----------------------------
+# Launch Gradio App
+# -----------------------------
 demo.launch()