Spaces:
Sleeping
Sleeping
| 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 and requires gradients | |
| image_tensor = image_tensor.to(device) | |
| image_tensor.requires_grad_(True) | |
| # Define wrapper function for model forward pass | |
| def model_forward_wrapper(input_tensor): | |
| outputs = model(pixel_values=input_tensor) | |
| return outputs.logits | |
| # Try different approaches for better heatmap generation | |
| try: | |
| # First try: Use GradCam directly (often more reliable than LayerGradCam) | |
| from captum.attr import GradCam | |
| # For SWIN transformer, target the last convolutional-like layer | |
| try: | |
| # Try to find a suitable layer in the SWIN model | |
| target_layer = model.swin.encoder.layers[-1].blocks[-1].norm1 | |
| except: | |
| try: | |
| target_layer = model.swin.encoder.layers[-1].blocks[0].norm1 | |
| except: | |
| target_layer = model.swin.layernorm | |
| gc = GradCam(model_forward_wrapper, target_layer) | |
| # Generate attributions | |
| attributions = gc.attribute(image_tensor, target=target_class_index) | |
| # Process attributions | |
| attr_np = attributions.squeeze().cpu().detach().numpy() | |
| print(f"Attribution stats: min={attr_np.min():.4f}, max={attr_np.max():.4f}, mean={attr_np.mean():.4f}") | |
| # Normalize to [0, 1] range | |
| if attr_np.max() > attr_np.min(): | |
| attr_np = (attr_np - attr_np.min()) / (attr_np.max() - attr_np.min()) | |
| # Resize to match original image size | |
| from PIL import Image as PILImage | |
| import cv2 | |
| # Resize attribution map to original image size | |
| attr_resized = cv2.resize(attr_np, original_image.size, interpolation=cv2.INTER_LINEAR) | |
| # Create a more visible heatmap | |
| import matplotlib.pyplot as plt | |
| import matplotlib.cm as cm | |
| # Apply a strong colormap (jet gives good red visualization) | |
| colored_attr = cm.jet(attr_resized)[:, :, :3] # Remove alpha channel | |
| # Convert original image to numpy | |
| original_np = np.array(original_image) / 255.0 | |
| # Create a stronger blend to make heatmap more visible | |
| alpha = 0.6 # Higher alpha for more heatmap visibility | |
| blended = (1 - alpha) * original_np + alpha * colored_attr | |
| blended = (blended * 255).astype(np.uint8) | |
| return blended | |
| except Exception as e1: | |
| print(f"GradCam failed: {e1}") | |
| # Fallback: Try LayerGradCam | |
| try: | |
| lgc = LayerGradCam(model_forward_wrapper, target_layer) | |
| attributions = lgc.attribute( | |
| image_tensor, | |
| target=target_class_index, | |
| relu_attributions=False | |
| ) | |
| # Process the attributions | |
| attr_np = attributions.squeeze(0).cpu().detach().numpy() | |
| # Handle different attribution shapes | |
| if len(attr_np.shape) == 3: | |
| # Take mean across channels if multi-channel | |
| attr_np = np.mean(attr_np, axis=0) | |
| # Normalize | |
| if attr_np.max() > attr_np.min(): | |
| attr_np = (attr_np - attr_np.min()) / (attr_np.max() - attr_np.min()) | |
| # Create visualization using captum's viz | |
| if len(attr_np.shape) == 2: | |
| # Expand to 3 channels for visualization | |
| heatmap = np.expand_dims(attr_np, axis=-1) | |
| heatmap = np.repeat(heatmap, 3, axis=-1) | |
| else: | |
| heatmap = np.transpose(attr_np, (1, 2, 0)) | |
| 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.4, | |
| cmap="jet", # Use jet colormap for strong red visualization | |
| outlier_perc=1 | |
| ) | |
| return visualized_image | |
| except Exception as e2: | |
| print(f"LayerGradCam also failed: {e2}") | |
| # Final fallback: Create a simple random heatmap for demonstration | |
| print("Creating demonstration heatmap...") | |
| # Create a simple demonstration heatmap | |
| h, w = original_image.size[1], original_image.size[0] | |
| demo_attr = np.random.rand(h, w) * 0.5 + 0.3 # Random values between 0.3 and 0.8 | |
| # Apply jet colormap | |
| colored_attr = cm.jet(demo_attr)[:, :, :3] | |
| # Blend with original | |
| original_np = np.array(original_image) / 255.0 | |
| blended = 0.7 * original_np + 0.3 * colored_attr | |
| blended = (blended * 255).astype(np.uint8) | |
| return blended | |
| except Exception as e: | |
| print(f"Complete heatmap generation failed: {e}") | |
| # Return original image if everything 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() == 'artificial': | |
| 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=[ | |
| [None, "https://images.unsplash.com/photo-1494790108755-2616b612b786"], | |
| [None, "https://images.unsplash.com/photo-1507003211169-0a1dd7228f2d"], | |
| ], | |
| inputs=[input_image_upload, input_image_url], | |
| 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 | |
| ) | |