app.py

# -*- coding: utf-8 -*-
import gradio as gr
import cv2
import numpy as np
import torch
from PIL import Image
from ultralytics import YOLO
from transformers import BlipProcessor, BlipForConditionalGeneration
from huggingface_hub import hf_hub_download
import warnings
warnings.filterwarnings('ignore')

# ============================================================
# 1. LOAD MODELS (CHANGE USERNAME BELOW)
# ============================================================
print("Loading YOLOv11 model from Hugging Face Hub...")
# ⚠️ IMPORTANT: Replace "YOUR_USERNAME" with your actual Hugging Face username
model_path = hf_hub_download(
    repo_id="SkyGuardAI/drone-detection-yolov11",  # <--- CHANGE THIS
    filename="best.pt"
)
model = YOLO(model_path)
print("YOLO model loaded successfully.")

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

print("Loading BLIP model...")
blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
blip_model.eval()
print("BLIP model loaded successfully.")

# ============================================================
# 2. HEATMAP FUNCTIONS (unchanged)
# ============================================================
layer_outputs = {}

def hook_fn(module, input, output):
    layer_outputs['feature_map'] = output.detach()

def get_best_layer(model):
    best_layer = None
    best_depth = 0
    pytorch_model = model.model if hasattr(model, 'model') else model
    for name, layer in pytorch_model.named_modules():
        if isinstance(layer, torch.nn.Conv2d):
            depth = name.count('.')
            if depth > best_depth:
                best_depth = depth
                best_layer = layer
    return best_layer

def generate_heatmap(model, image):
    try:
        layer_outputs.clear()
        if isinstance(image, Image.Image):
            image = np.array(image)
        img_resized = cv2.resize(image, (640, 640))
        pytorch_model = model.model if hasattr(model, 'model') else model
        target_layer = get_best_layer(pytorch_model)
        if target_layer is None:
            return None
        hook = target_layer.register_forward_hook(hook_fn)
        results = model(img_resized)
        hook.remove()
        if 'feature_map' not in layer_outputs:
            return None
        feature_map = layer_outputs['feature_map']
        if feature_map.dim() == 4:
            heatmap = feature_map[0].mean(dim=0).cpu().numpy()
        else:
            heatmap = feature_map.cpu().numpy()
        heatmap = cv2.GaussianBlur(heatmap, (5, 5), 0)
        min_val = heatmap.min()
        max_val = heatmap.max()
        if max_val - min_val > 1e-8:
            heatmap = (heatmap - min_val) / (max_val - min_val)
        else:
            heatmap = np.zeros_like(heatmap)
        heatmap = cv2.resize(heatmap, (640, 640))
        threshold = np.percentile(heatmap, 70)
        heatmap = np.where(heatmap > threshold, heatmap, 0)
        if heatmap.max() > 0:
            heatmap = heatmap / heatmap.max()
        heatmap_colored = cv2.applyColorMap((heatmap * 255).astype(np.uint8), cv2.COLORMAP_JET)
        heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
        overlay = cv2.addWeighted(img_resized, 0.6, heatmap_colored, 0.4, 0)
        return overlay
    except Exception as e:
        print(f"Heatmap error: {e}")
        return None

# ============================================================
# 3. DYNAMIC CAPTION (BLIP)
# ============================================================
def generate_dynamic_caption(image):
    try:
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image).convert('RGB')
        elif isinstance(image, Image.Image):
            image = image.convert('RGB')
        inputs = blip_processor(image, return_tensors="pt").to(device)
        with torch.no_grad():
            out = blip_model.generate(
                **inputs,
                max_length=60,
                num_beams=5,
                temperature=0.7,
                repetition_penalty=1.2
            )
        caption = blip_processor.decode(out[0], skip_special_tokens=True)
        return caption
    except Exception as e:
        print(f"Caption error: {e}")
        return "AI model is analyzing the scene."

# ============================================================
# 4. XAI REPORT
# ============================================================
def build_xai_report(is_drone, confidence, drone_count, processing_time, image_caption):
    if is_drone:
        drone_text = "a drone" if drone_count == 1 else f"{drone_count} drones"
        if confidence >= 0.8:
            confidence_level = "VERY HIGH"
            confidence_assessment = "excellent"
        elif confidence >= 0.6:
            confidence_level = "HIGH"
            confidence_assessment = "good"
        elif confidence >= 0.5:
            confidence_level = "MODERATE"
            confidence_assessment = "acceptable"
        else:
            confidence_level = "LOW"
            confidence_assessment = "limited"
        
        report = f"""
================================================================================
                         XAI DRONE DETECTION REPORT
================================================================================

[DYNAMIC IMAGE ANALYSIS]
{image_caption}

[DETECTION RESULTS]
• Status: CONFIRMED
• Confidence: {confidence:.1%} (Level: {confidence_level})
• Drones Detected: {drone_count}
• Processing Time: {processing_time:.0f} milliseconds
• Model: YOLOv11
• XAI Method: Convolutional Feature Map Extraction

[XAI HEATMAP INTERPRETATION]
The heatmap shows red regions where the neural network focused its attention.
Strong red activation on {drone_text} confirms the model successfully learned
discriminative features for drone detection.

The model demonstrates {confidence_assessment} confidence, as evidenced by
the concentrated activation pattern in the heatmap.

[TECHNICAL DETAILS]
• Heatmap: Extracted from deepest convolutional layer of YOLOv11
• Color Code: Red = High activation (model focus) | Blue = Low activation
• Processing: Gaussian blur (5x5 kernel)
• Threshold: Top 30% activation retained

[XAI CONCLUSION]
The YOLOv11 model has successfully detected {drone_text} in this image with
{confidence_assessment} confidence. The heatmap confirms correct feature
learning as the neural network focused on the drone's location.
"""
    else:
        report = f"""
================================================================================
                         XAI DRONE DETECTION REPORT
================================================================================

[DYNAMIC IMAGE ANALYSIS]
{image_caption}

[DETECTION RESULTS]
• Status: NOT CONFIRMED
• Highest Confidence: {confidence:.1%}
• Processing Time: {processing_time:.0f} milliseconds
• Model: YOLOv11
• XAI Method: Convolutional Feature Map Extraction

[XAI HEATMAP INTERPRETATION]
The heatmap shows scattered or unfocused activation patterns without strong
concentration on any specific region. This indicates the model did not identify
strong drone-like features in this image.

[POSSIBLE REASONS]
• No drone is present in the image
• Drone is too small or too far from the camera
• Poor lighting conditions reduce feature visibility
• Image blur or motion blur affects detection quality
• Drone is partially occluded by other objects

[RECOMMENDATIONS]
• Ensure adequate lighting in the scene
• Position the drone closer to the camera
• Use higher resolution images without motion blur
• Avoid cluttered backgrounds that may confuse the model

[TECHNICAL NOTE]
The heatmap was extracted from the deepest convolutional layer of YOLOv11.
Scattered activation pattern confirms absence of strong drone-like features.
"""
    return report

# ============================================================
# 5. MAIN PIPELINE FUNCTION
# ============================================================
def drone_detection_pipeline(input_image):
    try:
        if isinstance(input_image, Image.Image):
            img = np.array(input_image)
        else:
            img = input_image.copy()
        
        original_h, original_w = img.shape[:2]
        results = model(img)
        heatmap_overlay = generate_heatmap(model, img)
        
        is_drone = False
        confidence = 0.0
        drone_count = 0
        if results and len(results) > 0 and hasattr(results[0], 'boxes'):
            boxes_data = results[0].boxes
            if boxes_data and boxes_data.data is not None:
                data = boxes_data.data.cpu().numpy()
                for det in data:
                    if len(det) >= 6:
                        conf = float(det[4])
                        cls = int(det[5])
                        class_name = model.names[cls]
                        if class_name.lower() == 'drone' and conf >= 0.3:
                            drone_count += 1
                            confidence = max(confidence, conf)
        is_drone = drone_count > 0
        
        result_img = results[0].plot() if len(results) > 0 else img
        result_img_rgb = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
        caption = generate_dynamic_caption(img)
        processing_time_ms = 0
        report = build_xai_report(is_drone, confidence, drone_count, processing_time_ms, caption)
        
        if heatmap_overlay is not None:
            heatmap_resized = cv2.resize(heatmap_overlay, (original_w, original_h))
        else:
            heatmap_resized = np.zeros_like(result_img_rgb)
        
        return result_img_rgb, heatmap_resized, report
    except Exception as e:
        error_msg = f"An error occurred during processing: {str(e)}"
        print(error_msg)
        blank = np.zeros((480, 640, 3), dtype=np.uint8)
        return blank, blank, error_msg

# ============================================================
# 6. INTERFACE - DARK BLUE ONLY, NO PURPLE/VIOLET
# ============================================================
custom_css = """
/* Force white background */
.gradio-container, body, .gradio-container .main, .gradio-container .gradio-container {
    background: #ffffff !important;
    font-family: 'Segoe UI', 'Roboto', sans-serif;
}
/* Override any purple/violet colors from Gradio default theme */
:root {
    --primary-50: #0a2b4e !important;
    --primary-100: #1e4a76 !important;
    --primary-200: #2c5282 !important;
    --primary-300: #1e3a8a !important;
    --primary-400: #0a2b4e !important;
    --primary-500: #0a2b4e !important;
    --primary-600: #0a2b4e !important;
    --primary-700: #0a2b4e !important;
    --primary-800: #0a2b4e !important;
    --primary-900: #0a2b4e !important;
}
/* Header - dark blue */
.skyguard-header {
    text-align: center;
    padding: 1.5rem 0 0.5rem 0;
    margin-bottom: 1.5rem;
}
.skyguard-header h1 {
    font-size: 3rem;
    font-weight: 800;
    letter-spacing: 3px;
    color: #0a2b4e !important;
    margin-bottom: 0.2rem;
}
.skyguard-header .tagline {
    font-size: 1.2rem;
    font-weight: 600;
    color: #1e4a76 !important;
    letter-spacing: 1px;
    margin-top: 0.2rem;
}
.skyguard-header .sub {
    font-size: 0.9rem;
    color: #2c5282 !important;
    margin-top: 0.5rem;
}
/* Primary button - solid dark blue, no gradient, no purple */
.gr-button-primary, button.gr-button-primary, .gr-button.primary {
    background: #0a2b4e !important;
    border: none !important;
    color: white !important;
    font-weight: 700 !important;
    border-radius: 30px !important;
    padding: 10px 28px !important;
    transition: all 0.2s ease !important;
    text-transform: uppercase !important;
    letter-spacing: 1px;
    box-shadow: none !important;
}
.gr-button-primary:hover, button.gr-button-primary:hover, .gr-button.primary:hover {
    background: #1e4a76 !important;
    transform: scale(1.02);
    box-shadow: 0 2px 8px rgba(10,43,78,0.2) !important;
}
/* Tabs - remove purple */
.gr-tabs .tab-nav button {
    background-color: #eef2f5 !important;
    color: #0a2b4e !important;
    font-weight: 600 !important;
    border-radius: 8px 8px 0 0 !important;
    padding: 8px 20px !important;
    margin-right: 4px;
    border: none !important;
}
.gr-tabs .tab-nav button.selected {
    background-color: #0a2b4e !important;
    color: white !important;
}
/* Input and output boxes */
.gr-box, .gr-form, .gr-input, .gr-panel {
    background-color: #f9fafb !important;
    border: 1px solid #d1d5db !important;
    border-radius: 12px !important;
}
/* Labels */
label, .gr-form label {
    color: #1e293b !important;
    font-weight: 500 !important;
}
/* Markdown text */
.gr-markdown p, .gr-markdown {
    color: #1e293b !important;
}
/* Footer */
.skyguard-footer {
    text-align: center;
    margin-top: 30px;
    font-size: 12px;
    color: #6c757d;
    border-top: 1px solid #e2e8f0;
    padding-top: 15px;
}
/* Additional overrides for any purple elements */
.gr-button, .gr-button-lg, .gr-button-sm, button {
    --tw-ring-color: #0a2b4e !important;
}
.prose a, a {
    color: #1e4a76 !important;
}
input:focus, textarea:focus, select:focus {
    border-color: #0a2b4e !important;
    ring-color: #0a2b4e !important;
}
"""

with gr.Blocks(title="SkyGuard - Drone Detection System", theme=gr.themes.Soft(), css=custom_css) as demo:
    gr.HTML("""
        <div class="skyguard-header">
            <h1>SKYGUARD</h1>
            <div class="tagline">SEE EVERY THREAT BEFORE IT LANDS</div>
            <div class="sub">YOLOv11 Detection | Explainable AI (XAI) | BLIP Captioning</div>
            <div class="sub" style="font-size:0.75rem; color:#3b6e9e;">Mohamed Khider University of Biskra - Computer Science | Biskra, Algeria</div>
        </div>
    """)
    
    with gr.Row(equal_height=False):
        with gr.Column(scale=1, min_width=300):
            input_image = gr.Image(label="Upload Aerial Image", type="pil")
            analyze_btn = gr.Button("Analyze Threat", variant="primary")
        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.TabItem("Detection"):
                    output_image = gr.Image(label="Bounding Boxes")
                with gr.TabItem("Heatmap (XAI)"):
                    heatmap_image = gr.Image(label="Neural Activation Heatmap")
                with gr.TabItem("Intelligence Report"):
                    report_text = gr.Markdown(label="XAI Analysis Report")
    
    analyze_btn.click(
        fn=drone_detection_pipeline,
        inputs=[input_image],
        outputs=[output_image, heatmap_image, report_text]
    )
    
    gr.HTML("""
        <div class="skyguard-footer">
            SkyGuard System | Real-time Drone Detection | Powered by Hugging Face Spaces
        </div>
    """)

# ============================================================
# 7. RUN APP
# ============================================================
if __name__ == "__main__":
    demo.launch()