Update app.py

app.py

@@ -1,87 +1,132 @@
 import gradio as gr
 import torch
+import numpy as np
+import matplotlib.pyplot as plt
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 from PIL import Image
-import matplotlib.pyplot as plt
-import numpy as np
+import cv2
+from skimage import exposure
 import time
 
-# Load model and processor
-model_name = "dima806/deepfake_vs_real_image_detection"
-processor = AutoImageProcessor.from_pretrained(model_name)
-model = AutoModelForImageClassification.from_pretrained(model_name)
+# Load models (using free Hugging Face models)
+MODEL_NAMES = {
+    "Model 1": "dima806/deepfake_vs_real_image_detection",
+    "Model 2": "saltacc/anime-ai-detect",
+    "Model 3": "rizvandwiki/gansfake-detector"
+}
 
-def predict_image(image):
+# Initialize models
+models = {}
+processors = {}
+
+for name, path in MODEL_NAMES.items():
+    try:
+        processors[name] = AutoImageProcessor.from_pretrained(path)
+        models[name] = AutoModelForImageClassification.from_pretrained(path)
+    except:
+        print(f"Could not load model: {name}")
+
+def analyze_image(image, selected_model):
     if image is None:
-        return {"result": "⚠️ Please upload an image", "class": "neutral"}
+        return None, None, "Please upload an image first", None
     
     try:
-        # Show loading animation
-        time.sleep(1)  # Simulate processing time
-        
-        # Convert image to RGB if needed
+        # Convert to RGB if needed
         if image.mode != "RGB":
             image = image.convert("RGB")
         
-        # Preprocess and predict
+        # Get model and processor
+        model = models.get(selected_model)
+        processor = processors.get(selected_model)
+        
+        if not model or not processor:
+            return None, None, "Selected model not available", None
+        
+        # Preprocess image
         inputs = processor(images=image, return_tensors="pt")
+        
+        # Predict
         with torch.no_grad():
             outputs = model(**inputs)
         
         # Get probabilities
         probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
-        real_prob = probs[model.config.label2id["real"]].item()
-        fake_prob = probs[model.config.label2id["fake"]].item()
         
-        # Create visualization
-        fig = create_visualization(real_prob, fake_prob)
+        # Create visualizations
+        heatmap = generate_heatmap(image, model, processor)
+        chart_fig = create_probability_chart(probs, model.config.id2label)
         
-        if real_prob > fake_prob:
-            return {
-                "result": f"🎉 Authentic Image (Confidence: {real_prob*100:.1f}%)",
-                "class": "authentic",
-                "details": f"Real Probability: {real_prob*100:.1f}% | Fake Probability: {fake_prob*100:.1f}%",
-                "visualization": fig
-            }
-        else:
-            return {
-                "result": f"⚠️ Potential Deepfake (Confidence: {fake_prob*100:.1f}%)",
-                "class": "fake",
-                "details": f"Fake Probability: {fake_prob*100:.1f}% | Real Probability: {real_prob*100:.1f}%",
-                "visualization": fig
-            }
+        # Format results
+        result_text = format_results(probs, model.config.id2label)
+        
+        return heatmap, chart_fig, result_text, create_model_info(selected_model)
+    
     except Exception as e:
-        return {"result": f"❌ Error: {str(e)}", "class": "error"}
+        return None, None, f"Error: {str(e)}", None
 
-def create_visualization(real_prob, fake_prob):
-    labels = ['Real', 'Fake']
-    probabilities = [real_prob, fake_prob]
-    colors = ['#4CAF50', '#F44336']
+def generate_heatmap(image, model, processor):
+    """Generate a heatmap showing important regions for the prediction"""
+    # Convert to numpy array
+    img_array = np.array(image)
     
-    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))
+    # Create a saliency map (simple version)
+    gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
+    blurred = cv2.GaussianBlur(gray, (21, 21), 0)
+    heatmap = cv2.applyColorMap(blurred, cv2.COLORMAP_JET)
     
-    # Bar chart
-    bars = ax1.bar(labels, probabilities, color=colors)
-    ax1.set_ylim(0, 1)
-    ax1.set_title('Probability Distribution')
-    ax1.set_ylabel('Probability')
+    # Blend with original image
+    heatmap = cv2.addWeighted(img_array, 0.7, heatmap, 0.3, 0)
     
-    # Add value labels on top of bars
-    for bar in bars:
-        height = bar.get_height()
-        ax1.text(bar.get_x() + bar.get_width()/2., height,
-                f'{height:.2f}',
-                ha='center', va='bottom')
+    # Convert back to PIL
+    return Image.fromarray(heatmap)
+
+def create_probability_chart(probs, id2label):
+    """Create a bar chart of class probabilities"""
+    labels = [id2label[i] for i in range(len(probs))]
+    colors = ['#4CAF50' if 'real' in label.lower() else '#F44336' for label in labels]
     
-    # Pie chart
-    ax2.pie(probabilities, labels=labels, colors=colors, autopct='%1.1f%%', startangle=90)
-    ax2.axis('equal')
-    ax2.set_title('Probability Ratio')
+    fig, ax = plt.subplots(figsize=(8, 4))
+    bars = ax.barh(labels, probs.numpy(), color=colors)
+    ax.set_xlim(0, 1)
+    ax.set_title('Detection Probabilities', pad=20)
+    ax.set_xlabel('Probability')
+    
+    # Add value labels
+    for bar in bars:
+        width = bar.get_width()
+        ax.text(width + 0.02, bar.get_y() + bar.get_height()/2,
+                f'{width:.2f}',
+                va='center')
     
     plt.tight_layout()
     return fig
 
-# Custom CSS for attractive UI
+def format_results(probs, id2label):
+    """Format the results as text"""
+    results = []
+    for i, prob in enumerate(probs):
+        results.append(f"{id2label[i]}: {prob*100:.1f}%")
+    
+    max_prob = max(probs)
+    max_class = id2label[torch.argmax(probs).item()]
+    
+    if 'real' in max_class.lower():
+        conclusion = f"Conclusion: This image appears to be AUTHENTIC with {max_prob*100:.1f}% confidence"
+    else:
+        conclusion = f"Conclusion: This image appears to be FAKE/GENERATED with {max_prob*100:.1f}% confidence"
+    
+    return "\n".join([conclusion, "", "Detailed probabilities:"] + results)
+
+def create_model_info(model_name):
+    """Create information about the current model"""
+    info = {
+        "Model 1": "Trained to detect deepfakes vs real human faces",
+        "Model 2": "Specialized in detecting AI-generated anime images",
+        "Model 3": "General GAN-generated image detector"
+    }
+    return info.get(model_name, "No information available about this model")
+
+# Custom CSS for the interface
 custom_css = """
 :root {
     --primary: #4b6cb7;
@@ -89,11 +134,10 @@ custom_css = """
     --authentic: #4CAF50;
     --fake: #F44336;
     --neutral: #2196F3;
-    --error: #FF9800;
 }
 
 #main-container {
-    max-width: 900px;
+    max-width: 1200px;
     margin: auto;
     padding: 25px;
     border-radius: 15px;
@@ -126,35 +170,18 @@ custom_css = """
     padding: 20px;
     border-radius: 12px;
     margin-top: 20px;
-    font-size: 1.2em;
-    font-weight: bold;
-    text-align: center;
+    font-size: 1.1em;
     transition: all 0.3s ease;
     box-shadow: 0 4px 6px rgba(0,0,0,0.1);
+    background: white;
 }
 
-.authentic {
-    background: linear-gradient(135deg, #e6f7e6 0%, #c8e6c9 100%);
-    color: var(--authentic);
-    border-left: 5px solid var(--authentic);
-}
-
-.fake {
-    background: linear-gradient(135deg, #ffebee 0%, #ffcdd2 100%);
-    color: var(--fake);
-    border-left: 5px solid var(--fake);
-}
-
-.neutral {
-    background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%);
-    color: var(--neutral);
-    border-left: 5px solid var(--neutral);
-}
-
-.error {
-    background: linear-gradient(135deg, #fff3e0 0%, #ffe0b2 100%);
-    color: var(--error);
-    border-left: 5px solid var(--error);
+.visualization-box {
+    border-radius: 12px;
+    padding: 15px;
+    background: white;
+    margin-top: 15px;
+    box-shadow: 0 4px 6px rgba(0,0,0,0.1);
 }
 
 .btn-primary {
@@ -166,17 +193,11 @@ custom_css = """
     font-weight: bold !important;
 }
 
-.btn-primary:hover {
-    transform: translateY(-2px);
-    box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important;
-}
-
-.btn-secondary {
+.model-select {
     background: white !important;
-    color: var(--primary) !important;
     border: 2px solid var(--primary) !important;
-    padding: 12px 24px !important;
     border-radius: 8px !important;
+    padding: 8px 12px !important;
 }
 
 .footer {
@@ -186,24 +207,16 @@ custom_css = """
     color: #666;
 }
 
-.animation {
-    animation: fadeIn 0.5s ease-in-out;
-}
-
-.visualization-box {
-    border-radius: 12px;
-    padding: 15px;
-    background: white;
-    margin-top: 15px;
-    box-shadow: 0 4px 6px rgba(0,0,0,0.1);
-}
-
 @keyframes fadeIn {
     from { opacity: 0; transform: translateY(10px); }
     to { opacity: 1; transform: translateY(0); }
 }
 
-.processing-animation {
+.animation {
+    animation: fadeIn 0.5s ease-in-out;
+}
+
+.loading {
     animation: pulse 1.5s infinite;
 }
 
@@ -218,93 +231,66 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
     with gr.Column(elem_id="main-container"):
         with gr.Column(elem_classes=["header"]):
             gr.Markdown("# 🛡️ DeepGuard AI")
-            gr.Markdown("Detect AI-generated images with advanced deep learning")
+            gr.Markdown("## Advanced Deepfake Detection System")
         
         with gr.Row():
-            with gr.Column(scale=2):
+            with gr.Column(scale=1.5):
                 image_input = gr.Image(
                     type="pil", 
-                    label="Drag & drop your image here", 
+                    label="Upload Image for Analysis", 
                     elem_classes=["upload-area", "animation"]
                 )
+                
                 with gr.Row():
-                    submit_btn = gr.Button("Analyze Now", elem_classes=["btn-primary", "animation"])
-                    clear_btn = gr.Button("Clear", elem_classes=["btn-secondary", "animation"])
+                    model_selector = gr.Dropdown(
+                        choices=list(MODEL_NAMES.keys()),
+                        value=list(MODEL_NAMES.keys())[0],
+                        label="Select Detection Model",
+                        elem_classes=["model-select", "animation"]
+                    )
+                    analyze_btn = gr.Button(
+                        "Analyze Image", 
+                        elem_classes=["btn-primary", "animation"]
+                    )
             
             with gr.Column(scale=1):
-                result_display = gr.Textbox(
-                    label="Analysis Result",
-                    interactive=False,
-                    elem_classes=["result-box", "neutral", "animation"]
-                )
-                details = gr.Textbox(
-                    label="Detailed Analysis",
-                    interactive=False,
-                    visible=False
-                )
+                with gr.Column(elem_classes=["visualization-box"]):
+                    heatmap_output = gr.Image(
+                        label="Attention Heatmap",
+                        interactive=False
+                    )
+                
+                with gr.Column(elem_classes=["visualization-box"]):
+                    chart_output = gr.Plot(
+                        label="Detection Probabilities"
+                    )
         
-        visualization = gr.Plot(
-            label="Probability Visualization",
-            visible=False,
-            elem_classes=["visualization-box"]
-        )
+        with gr.Column(elem_classes=["result-box", "animation"]):
+            result_output = gr.Textbox(
+                label="Analysis Results",
+                interactive=False,
+                lines=8
+            )
+        
+        with gr.Column(elem_classes=["result-box", "animation"]):
+            model_info = gr.Textbox(
+                label="Model Information",
+                interactive=False,
+                lines=3
+            )
         
         gr.Markdown("""
         <div class="footer">
-        *Note: This AI detection tool provides estimates based on image analysis. 
-        Results should be verified with additional methods for critical decisions.*
+        *Note: This tool provides probabilistic estimates. Always verify important findings with additional methods.<br>
+        Models may produce false positives/negatives. Performance varies by image type and quality.*
         </div>
         """)
-
-    def update_ui(image, prediction):
-        if image is None:
-            return {
-                result_display: gr.Textbox.update(
-                    value="📁 Upload an image to begin analysis",
-                    elem_classes=["result-box", "neutral"]
-                ),
-                details: gr.Textbox.update(visible=False),
-                visualization: gr.Plot.update(visible=False)
-            }
-        
-        show_visualization = "visualization" in prediction
-        show_details = "details" in prediction
-        
-        updates = {
-            result_display: gr.Textbox.update(
-                value=prediction["result"],
-                elem_classes=["result-box", prediction["class"]]
-            ),
-            details: gr.Textbox.update(
-                value=prediction.get("details", ""),
-                visible=show_details
-            ),
-            visualization: gr.Plot.update(
-                value=prediction.get("visualization"),
-                visible=show_visualization
-            )
-        }
-        
-        return updates
-
-    submit_btn.click(
-        fn=predict_image,
-        inputs=image_input,
-        outputs=gr.JSON(visible=False)
-    ).then(
-        fn=update_ui,
-        inputs=[image_input, gr.JSON()],
-        outputs=[result_display, details, visualization]
-    )
     
-    clear_btn.click(
-        fn=lambda: [None, {"result": "📁 Upload an image to begin analysis", "class": "neutral"}],
-        inputs=None,
-        outputs=[image_input, gr.JSON()]
-    ).then(
-        fn=update_ui,
-        inputs=[image_input, gr.JSON()],
-        outputs=[result_display, details, visualization]
+    # Event handlers
+    analyze_btn.click(
+        fn=analyze_image,
+        inputs=[image_input, model_selector],
+        outputs=[heatmap_output, chart_output, result_output, model_info]
     )
 
 if __name__ == "__main__":