Create app.py

app.py

@@ -0,0 +1,274 @@
+import gradio as gr
+import torch
+from transformers import AutoModelForImageClassification, AutoImageProcessor
+from PIL import Image
+import numpy as np
+from captum.attr import LayerGradCam
+from captum.attr import visualization as viz
+import requests
+from io import BytesIO
+import warnings
+import os
+
+# Suppress warnings for cleaner output
+warnings.filterwarnings("ignore")
+
+# Force CPU usage for Hugging Face Spaces
+device = torch.device("cpu")
+torch.set_num_threads(1)  # Optimize for CPU usage
+
+# --- 1. Load Model and Processor ---
+print("Loading model and processor...")
+try:
+    model_id = "Organika/sdxl-detector"
+    processor = AutoImageProcessor.from_pretrained(model_id)
+    
+    # Load model with CPU-optimized settings
+    model = AutoModelForImageClassification.from_pretrained(
+        model_id, 
+        torch_dtype=torch.float32,
+        device_map="cpu",
+        low_cpu_mem_usage=True
+    )
+    model.to(device)
+    model.eval()
+    print("Model and processor loaded successfully on CPU.")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    raise
+
+# --- 2. Define the Explainability (Grad-CAM) Function ---
+def generate_heatmap(image_tensor, original_image, target_class_index):
+    try:
+        # Ensure tensor is on CPU
+        image_tensor = image_tensor.to(device)
+        
+        # Define wrapper function for model forward pass
+        def model_forward_wrapper(input_tensor):
+            with torch.no_grad():  # Save memory during attribution
+                outputs = model(pixel_values=input_tensor)
+                return outputs.logits
+
+        # Get the target layer for Grad-CAM
+        # For SWIN transformer, use the layer normalization layer
+        target_layer = model.swin.layernorm
+
+        # Initialize LayerGradCam with the wrapper function
+        lgc = LayerGradCam(model_forward_wrapper, target_layer)
+
+        # Generate attributions
+        with torch.no_grad():
+            attributions = lgc.attribute(
+                image_tensor, 
+                target=target_class_index, 
+                relu_attributions=True
+            )
+
+        # Convert attributions to numpy for visualization
+        heatmap = np.transpose(
+            attributions.squeeze(0).cpu().detach().numpy(), 
+            (1, 2, 0)
+        )
+
+        # Create visualization
+        visualized_image, _ = viz.visualize_image_attr(
+            heatmap,
+            np.array(original_image),
+            method="blended_heat_map",
+            sign="all",
+            show_colorbar=True,
+            title="AI Detection Heatmap",
+            alpha_overlay=0.6
+        )
+        
+        return visualized_image
+        
+    except Exception as e:
+        print(f"Error generating heatmap: {e}")
+        # Return original image if heatmap generation fails
+        return np.array(original_image)
+
+# --- 3. Main Prediction Function ---
+def predict(image_upload: Image.Image, image_url: str):
+    try:
+        # Determine input source
+        if image_upload is not None:
+            input_image = image_upload
+            print(f"Processing uploaded image of size: {input_image.size}")
+        elif image_url and image_url.strip():
+            try:
+                response = requests.get(image_url, timeout=10)
+                response.raise_for_status()
+                input_image = Image.open(BytesIO(response.content))
+                print(f"Processing image from URL: {image_url}")
+            except Exception as e:
+                raise gr.Error(f"Could not load image from URL. Please check the link. Error: {e}")
+        else:
+            raise gr.Error("Please upload an image or provide a URL to analyze.")
+
+        # Convert RGBA to RGB if necessary
+        if input_image.mode == 'RGBA':
+            input_image = input_image.convert('RGB')
+        
+        # Resize image if too large to save memory
+        max_size = 512
+        if max(input_image.size) > max_size:
+            input_image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
+
+        # Process image
+        inputs = processor(images=input_image, return_tensors="pt")
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+
+        # Make prediction
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits
+
+        # Calculate probabilities
+        probabilities = torch.nn.functional.softmax(logits, dim=-1)
+        predicted_class_idx = logits.argmax(-1).item()
+        confidence_score = probabilities[0][predicted_class_idx].item()
+        predicted_label = model.config.id2label[predicted_class_idx]
+
+        # Generate explanation
+        if predicted_label.lower() == 'ai':
+            explanation = (
+                f"🤖 The model is {confidence_score:.2%} confident that this image is **AI-GENERATED**.\n\n"
+                "The heatmap highlights areas that most influenced this decision. "
+                "Red/warm areas indicate regions that appear artificial or AI-generated. "
+                "Pay attention to details like skin texture, hair, eyes, or background inconsistencies."
+            )
+        else:
+            explanation = (
+                f"👤 The model is {confidence_score:.2%} confident that this image is **HUMAN-MADE**.\n\n"
+                "The heatmap shows areas the model considers natural and realistic. "
+                "Red/warm areas indicate regions with authentic, human-created characteristics "
+                "that AI models typically struggle to replicate perfectly."
+            )
+
+        print("Generating heatmap...")
+        heatmap_image = generate_heatmap(inputs['pixel_values'], input_image, predicted_class_idx)
+        print("Heatmap generated successfully.")
+
+        # Create labels dictionary for gradio output
+        labels_dict = {
+            model.config.id2label[i]: float(probabilities[0][i]) 
+            for i in range(len(model.config.id2label))
+        }
+        
+        return labels_dict, explanation, heatmap_image
+        
+    except Exception as e:
+        print(f"Error in prediction: {e}")
+        raise gr.Error(f"An error occurred during prediction: {str(e)}")
+
+# --- 4. Gradio Interface ---
+with gr.Blocks(
+    theme=gr.themes.Soft(),
+    title="AI Image Detector",
+    css="""
+    .gradio-container {
+        max-width: 1200px !important;
+    }
+    .tab-nav {
+        margin-bottom: 1rem;
+    }
+    """
+) as demo:
+    gr.Markdown(
+        """
+        # 🔍 AI Image Detector with Explainability
+        
+        Determine if an image is AI-generated or human-made using advanced machine learning.
+        
+        **Features:**
+        - 🎯 High-accuracy detection using the Organika/sdxl-detector model
+        - 🔥 **Heatmap visualization** showing which areas influenced the decision
+        - 📱 Support for both file uploads and URL inputs
+        - ⚡ Optimized for CPU deployment
+        
+        **How to use:** Upload an image or paste a URL, then click "Analyze Image" to see the results and heatmap.
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### 📥 Input")
+            
+            with gr.Tabs():
+                with gr.TabItem("📁 Upload File"):
+                    input_image_upload = gr.Image(
+                        type="pil", 
+                        label="Upload Your Image",
+                        height=300
+                    )
+                with gr.TabItem("🔗 Use URL"):
+                    input_image_url = gr.Textbox(
+                        label="Paste Image URL here",
+                        placeholder="https://example.com/image.jpg"
+                    )
+
+            submit_btn = gr.Button(
+                "🔍 Analyze Image", 
+                variant="primary",
+                size="lg"
+            )
+            
+            gr.Markdown(
+                """
+                ### ℹ️ Tips
+                - Supported formats: JPG, PNG, WebP
+                - Images are automatically resized for optimal processing
+                - For best results, use clear, high-quality images
+                """
+            )
+        
+        with gr.Column(scale=2):
+            gr.Markdown("### 📊 Results")
+            
+            with gr.Row():
+                with gr.Column():
+                    output_label = gr.Label(
+                        label="Prediction Confidence",
+                        num_top_classes=2
+                    )
+                with gr.Column():
+                    output_text = gr.Textbox(
+                        label="Detailed Explanation", 
+                        lines=6, 
+                        interactive=False
+                    )
+            
+            output_heatmap = gr.Image(
+                label="🔥 AI Detection Heatmap - Red areas influenced the decision most",
+                height=400
+            )
+
+    # Connect the interface
+    submit_btn.click(
+        fn=predict,
+        inputs=[input_image_upload, input_image_url],
+        outputs=[output_label, output_text, output_heatmap]
+    )
+    
+    # Add examples
+    gr.Examples(
+        examples=[
+            ["https://images.unsplash.com/photo-1494790108755-2616b612b786", ""],
+            ["https://images.unsplash.com/photo-1507003211169-0a1dd7228f2d", ""],
+        ],
+        inputs=[input_image_url, input_image_upload],
+        outputs=[output_label, output_text, output_heatmap],
+        fn=predict,
+        cache_examples=False
+    )
+
+# --- 5. Launch the App ---
+if __name__ == "__main__":
+    demo.launch(
+        debug=False,
+        share=False,
+        server_name="0.0.0.0",
+        server_port=7860
+    )
+