app.py

# ═══════════════════════════════════════════════════════════════════
# app.py - Pothole Detection Gradio Application
# ═══════════════════════════════════════════════════════════════════

import gradio as gr
import cv2
import numpy as np
from PIL import Image
import torch
from ultralytics import YOLO
from pathlib import Path
import tempfile
import os

# ═══════════════════════════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════════════════════════

MODEL_PATH = "best.pt"  # Place your model in the same directory
CONFIDENCE_THRESHOLD = 0.5

# ═══════════════════════════════════════════════════════════════════
# LOAD MODEL
# ═══════════════════════════════════════════════════════════════════

print("🔄 Loading pothole detection model...")

if not Path(MODEL_PATH).exists():
    raise FileNotFoundError(f"❌ Model not found: {MODEL_PATH}\nPlease place your best.pt file in the app directory!")

model = YOLO(MODEL_PATH)
print("✅ Model loaded successfully!")

# ═══════════════════════════════════════════════════════════════════
# MEASUREMENT CALCULATOR
# ═══════════════════════════════════════════════════════════════════

class PotholeMeasurementSystem:
    """Calculate physical measurements from segmentation masks"""
    
    def __init__(self):
        # Intel RealSense D415 intrinsics (default camera calibration)
        self.fx = 384.0
        self.fy = 384.0
        self.cx = 320.0
        self.cy = 240.0
    
    def calculate_measurements(self, mask, depth_map=None):
        """
        Calculate all physical measurements from mask
        
        Args:
            mask: Binary segmentation mask (H, W)
            depth_map: Optional depth map for accurate measurements
        
        Returns:
            Dictionary with all measurements
        """
        mask_bool = mask > 0
        
        if not np.any(mask_bool):
            return None
        
        # Calculate perimeter
        contours, _ = cv2.findContours(
            mask.astype(np.uint8), 
            cv2.RETR_EXTERNAL, 
            cv2.CHAIN_APPROX_NONE
        )
        
        if len(contours) == 0:
            return None
        
        contour = max(contours, key=cv2.contourArea)
        
        # Pixel-based measurements (if no depth map available)
        pixel_area = np.sum(mask_bool)
        perimeter_pixels = cv2.arcLength(contour, True)
        
        # Estimate physical measurements (assumes camera at ~1m height)
        # Approximate scale: 1 pixel ≈ 0.5mm at 1m distance
        pixel_to_mm = 0.5  # This is an approximation
        
        if depth_map is not None:
            # Use depth map for accurate measurements
            h_c = np.median(depth_map[~mask_bool & (depth_map > 0)])
            pothole_depths = depth_map[mask_bool]
            pothole_depths = pothole_depths[pothole_depths > 0]
            
            if len(pothole_depths) > 0:
                actual_depths = pothole_depths - h_c
                positive_depths = actual_depths[actual_depths > 0]
                
                max_depth_mm = float(actual_depths.max()) if len(actual_depths) > 0 else 0
                mean_depth_mm = float(actual_depths.mean()) if len(actual_depths) > 0 else 0
                
                # Scale factors based on depth
                depth_m = h_c / 1000.0
                s_x = depth_m / self.fx
                s_y = depth_m / self.fy
                pixel_to_mm = (s_x + s_y) / 2 * 1000
            else:
                max_depth_mm = 0
                mean_depth_mm = 0
        else:
            # Estimate depth from area (heuristic)
            # Larger potholes tend to be deeper
            estimated_depth_cm = min(15, (pixel_area / 1000) * 2)
            max_depth_mm = estimated_depth_cm * 10
            mean_depth_mm = max_depth_mm * 0.7
        
        # Calculate physical dimensions
        perimeter_cm = (perimeter_pixels * pixel_to_mm) / 10
        area_cm2 = (pixel_area * pixel_to_mm * pixel_to_mm) / 100
        area_m2 = area_cm2 / 10000
        
        # Estimate volume (assuming cone-like shape)
        volume_liters = (area_m2 * (max_depth_mm / 1000) / 3) * 1000
        
        # Severity classification
        depth_cm = max_depth_mm / 10
        if depth_cm > 10 or area_m2 > 0.5:
            severity = 'CRITICAL'
            severity_color = '🔴'
        elif depth_cm > 5 or area_m2 > 0.2:
            severity = 'HIGH'
            severity_color = '🟠'
        elif depth_cm > 3:
            severity = 'MEDIUM'
            severity_color = '🟡'
        else:
            severity = 'LOW'
            severity_color = '🟢'
        
        return {
            'max_depth_mm': max_depth_mm,
            'max_depth_cm': max_depth_mm / 10,
            'mean_depth_mm': mean_depth_mm,
            'mean_depth_cm': mean_depth_mm / 10,
            'perimeter_cm': perimeter_cm,
            'perimeter_m': perimeter_cm / 100,
            'area_cm2': area_cm2,
            'area_m2': area_m2,
            'volume_liters': volume_liters,
            'volume_m3': volume_liters / 1000,
            'num_pixels': int(pixel_area),
            'severity': severity,
            'severity_color': severity_color,
            'contour': contour
        }

measurer = PotholeMeasurementSystem()

# ═══════════════════════════════════════════════════════════════════
# INFERENCE FUNCTION
# ═══════════════════════════════════════════════════════════════════

def detect_potholes(image, confidence_threshold=0.5, depth_map=None):
    """
    Main inference function for Gradio
    
    Args:
        image: PIL Image or numpy array
        confidence_threshold: Detection confidence threshold
        depth_map: Optional depth map (not used in basic deployment)
    
    Returns:
        annotated_image: Image with detections
        metrics_html: HTML formatted metrics
        summary_text: Text summary
    """
    
    # Convert PIL to numpy if needed
    if isinstance(image, Image.Image):
        image = np.array(image)
    
    # Ensure RGB format
    if len(image.shape) == 2:  # Grayscale
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    elif image.shape[2] == 4:  # RGBA
        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
    
    h, w = image.shape[:2]
    
    # Save to temporary file (YOLO needs file path)
    with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file:
        tmp_path = tmp_file.name
        cv2.imwrite(tmp_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
    
    try:
        # Run prediction
        results = model(tmp_path, conf=confidence_threshold, verbose=False)[0]
        
        # Check if any detections
        if results.boxes is None or len(results.boxes) == 0:
            return (
                image, 
                "<h3 style='color: orange;'>⚠️ No potholes detected</h3>",
                "No potholes detected in the image."
            )
        
        # Extract results
        boxes = results.boxes.xyxy.cpu().numpy()
        confidences = results.boxes.conf.cpu().numpy()
        masks = results.masks.data.cpu().numpy() if results.masks is not None else None
        
        # Create annotated image
        annotated_img = image.copy()
        
        # Store all measurements
        all_measurements = []
        
        # Process each detection
        for idx, (box, conf) in enumerate(zip(boxes, confidences)):
            x1, y1, x2, y2 = box.astype(int)
            
            # Draw bounding box
            cv2.rectangle(annotated_img, (x1, y1), (x2, y2), (255, 0, 0), 3)
            
            # Process mask if available
            if masks is not None and idx < len(masks):
                mask = masks[idx]
                mask_resized = cv2.resize(mask, (w, h))
                mask_binary = (mask_resized > 0.5).astype(np.uint8) * 255
                
                # Create colored overlay
                overlay = annotated_img.copy()
                overlay[mask_binary > 0] = [255, 50, 50]  # Red overlay
                annotated_img = cv2.addWeighted(annotated_img, 0.6, overlay, 0.4, 0)
                
                # Draw contour
                contours, _ = cv2.findContours(
                    mask_binary, 
                    cv2.RETR_EXTERNAL, 
                    cv2.CHAIN_APPROX_SIMPLE
                )
                cv2.drawContours(annotated_img, contours, -1, (0, 255, 0), 2)
                
                # Calculate measurements
                measurements = measurer.calculate_measurements(mask_binary, depth_map)
                
                if measurements:
                    measurements['pothole_id'] = idx + 1
                    measurements['confidence'] = float(conf)
                    all_measurements.append(measurements)
                    
                    # Add text annotation
                    severity_emoji = measurements['severity_color']
                    text = f"#{idx+1} {severity_emoji} {measurements['severity']}"
                    
                    # Text background
                    text_size = cv2.getTextSize(
                        text, 
                        cv2.FONT_HERSHEY_SIMPLEX, 
                        0.7, 
                        2
                    )[0]
                    
                    cv2.rectangle(
                        annotated_img,
                        (x1, y1 - text_size[1] - 10),
                        (x1 + text_size[0] + 10, y1),
                        (0, 0, 0),
                        -1
                    )
                    
                    cv2.putText(
                        annotated_img, 
                        text, 
                        (x1 + 5, y1 - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 
                        0.7, 
                        (255, 255, 255), 
                        2
                    )
        
        # Generate metrics HTML
        metrics_html = generate_metrics_html(all_measurements)
        
        # Generate summary text
        summary_text = generate_summary_text(all_measurements)
        
        return annotated_img, metrics_html, summary_text
    
    finally:
        # Clean up temporary file
        if os.path.exists(tmp_path):
            os.unlink(tmp_path)

# ═══════════════════════════════════════════════════════════════════
# HTML GENERATION FUNCTIONS
# ═══════════════════════════════════════════════════════════════════

def generate_metrics_html(measurements):
    """Generate beautiful HTML metrics table"""
    
    if not measurements:
        return "<h3 style='color: orange;'>No measurements available</h3>"
    
    html = """
    <style>
        .metrics-container {
            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
            padding: 20px;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            border-radius: 15px;
            color: white;
        }
        .metrics-header {
            text-align: center;
            margin-bottom: 20px;
            font-size: 24px;
            font-weight: bold;
            text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
        }
        .pothole-card {
            background: rgba(255, 255, 255, 0.95);
            border-radius: 10px;
            padding: 15px;
            margin: 10px 0;
            color: #333;
            box-shadow: 0 4px 6px rgba(0,0,0,0.1);
        }
        .pothole-header {
            font-size: 18px;
            font-weight: bold;
            margin-bottom: 10px;
            display: flex;
            justify-content: space-between;
            align-items: center;
        }
        .severity-badge {
            padding: 5px 15px;
            border-radius: 20px;
            font-weight: bold;
            font-size: 14px;
        }
        .severity-CRITICAL { background: #ff4444; color: white; }
        .severity-HIGH { background: #ff9800; color: white; }
        .severity-MEDIUM { background: #ffeb3b; color: #333; }
        .severity-LOW { background: #4caf50; color: white; }
        .metrics-grid {
            display: grid;
            grid-template-columns: repeat(2, 1fr);
            gap: 10px;
            margin-top: 10px;
        }
        .metric-item {
            background: #f5f5f5;
            padding: 10px;
            border-radius: 5px;
            border-left: 4px solid #667eea;
        }
        .metric-label {
            font-size: 12px;
            color: #666;
            margin-bottom: 3px;
        }
        .metric-value {
            font-size: 18px;
            font-weight: bold;
            color: #333;
        }
        .summary-section {
            margin-top: 20px;
            padding: 15px;
            background: rgba(255, 255, 255, 0.95);
            border-radius: 10px;
            color: #333;
        }
        .summary-title {
            font-size: 18px;
            font-weight: bold;
            margin-bottom: 10px;
        }
    </style>
    
    <div class="metrics-container">
        <div class="metrics-header">
            🕳️ Pothole Detection Results
        </div>
    """
    
    # Individual pothole cards
    for m in measurements:
        severity_class = f"severity-{m['severity']}"
        
        html += f"""
        <div class="pothole-card">
            <div class="pothole-header">
                <span>{m['severity_color']} Pothole #{m['pothole_id']}</span>
                <span class="severity-badge {severity_class}">{m['severity']}</span>
            </div>
            
            <div style="margin-bottom: 10px;">
                <strong>Confidence:</strong> {m['confidence']*100:.1f}%
            </div>
            
            <div class="metrics-grid">
                <div class="metric-item">
                    <div class="metric-label">📏 Max Depth</div>
                    <div class="metric-value">{m['max_depth_cm']:.2f} cm</div>
                </div>
                
                <div class="metric-item">
                    <div class="metric-label">📐 Mean Depth</div>
                    <div class="metric-value">{m['mean_depth_cm']:.2f} cm</div>
                </div>
                
                <div class="metric-item">
                    <div class="metric-label">⭕ Perimeter</div>
                    <div class="metric-value">{m['perimeter_cm']:.2f} cm</div>
                </div>
                
                <div class="metric-item">
                    <div class="metric-label">📦 Area</div>
                    <div class="metric-value">{m['area_m2']:.4f} m²</div>
                </div>
                
                <div class="metric-item">
                    <div class="metric-label">💧 Volume</div>
                    <div class="metric-value">{m['volume_liters']:.2f} L</div>
                </div>
                
                <div class="metric-item">
                    <div class="metric-label">🔢 Pixels</div>
                    <div class="metric-value">{m['num_pixels']:,}</div>
                </div>
            </div>
        </div>
        """
    
    # Summary section
    total_area = sum(m['area_m2'] for m in measurements)
    total_volume = sum(m['volume_liters'] for m in measurements)
    avg_depth = np.mean([m['max_depth_cm'] for m in measurements])
    max_depth = max(m['max_depth_cm'] for m in measurements)
    
    severity_counts = {}
    for m in measurements:
        severity_counts[m['severity']] = severity_counts.get(m['severity'], 0) + 1
    
    html += f"""
        <div class="summary-section">
            <div class="summary-title">📊 Overall Summary</div>
            <div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 10px;">
                <div><strong>Total Potholes:</strong> {len(measurements)}</div>
                <div><strong>Average Depth:</strong> {avg_depth:.2f} cm</div>
                <div><strong>Maximum Depth:</strong> {max_depth:.2f} cm</div>
                <div><strong>Total Area:</strong> {total_area:.4f} m²</div>
                <div><strong>Total Volume:</strong> {total_volume:.2f} L</div>
                <div><strong>Severity Distribution:</strong> {', '.join([f"{k}: {v}" for k, v in severity_counts.items()])}</div>
            </div>
        </div>
    </div>
    """
    
    return html

def generate_summary_text(measurements):
    """Generate plain text summary"""
    
    if not measurements:
        return "No potholes detected."
    
    summary = f"🔍 DETECTION SUMMARY\n"
    summary += f"{'='*50}\n\n"
    
    summary += f"Total Potholes Detected: {len(measurements)}\n\n"
    
    for m in measurements:
        summary += f"{m['severity_color']} Pothole #{m['pothole_id']} - {m['severity']}\n"
        summary += f"   Confidence: {m['confidence']*100:.1f}%\n"
        summary += f"   Depth: {m['max_depth_cm']:.2f} cm\n"
        summary += f"   Area: {m['area_m2']:.4f} m²\n"
        summary += f"   Volume: {m['volume_liters']:.2f} liters\n\n"
    
    # Overall stats
    summary += f"{'='*50}\n"
    summary += f"OVERALL STATISTICS\n"
    summary += f"{'='*50}\n"
    summary += f"Average Depth: {np.mean([m['max_depth_cm'] for m in measurements]):.2f} cm\n"
    summary += f"Total Area: {sum(m['area_m2'] for m in measurements):.4f} m²\n"
    summary += f"Total Volume: {sum(m['volume_liters'] for m in measurements):.2f} liters\n"
    
    return summary

# ═══════════════════════════════════════════════════════════════════
# GRADIO INTERFACE
# ═══════════════════════════════════════════════════════════════════

# Custom CSS
custom_css = """
.gradio-container {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
}

#title {
    text-align: center;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    padding: 30px;
    border-radius: 10px;
    margin-bottom: 20px;
}

#subtitle {
    text-align: center;
    color: #666;
    margin-bottom: 30px;
}

.image-container {
    border: 3px solid #667eea;
    border-radius: 10px;
    padding: 10px;
}
"""

# Create Gradio interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as app:
    
    # Header
    gr.HTML("""
        <div id="title">
            <h1>🕳️ AI-Powered Pothole Detection System</h1>
            <p style="font-size: 18px; margin-top: 10px;">
                Advanced YOLOv11n Segmentation with Physical Measurements
            </p>
        </div>
    """)
    
    gr.HTML("""
        <div id="subtitle">
            <p style="font-size: 16px;">
                Upload an image to detect and measure potholes with high accuracy.<br>
                Get detailed metrics including depth, area, volume, and severity classification.
            </p>
        </div>
    """)
    
    # Main interface
    with gr.Row():
        with gr.Column(scale=1):
            # Input section
            gr.Markdown("### 📤 Upload Image")
            input_image = gr.Image(
                label="Road Image",
                type="pil",
                height=400
            )
            
            confidence_slider = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.5,
                step=0.05,
                label="🎯 Confidence Threshold",
                info="Lower = more detections, Higher = fewer but more confident"
            )
            
            detect_button = gr.Button(
                "🔍 Detect Potholes",
                variant="primary",
                size="lg"
            )
            
            gr.Markdown("""
            ---
            ### 📖 Instructions
            1. Upload a road image
            2. Adjust confidence threshold if needed
            3. Click "Detect Potholes"
            4. View results and metrics
            
            ### 📊 Severity Levels
            - 🔴 **CRITICAL**: Depth > 10cm or Area > 0.5m²
            - 🟠 **HIGH**: Depth > 5cm or Area > 0.2m²
            - 🟡 **MEDIUM**: Depth > 3cm
            - 🟢 **LOW**: Depth < 3cm
            """)
        
        with gr.Column(scale=1):
            # Output section
            gr.Markdown("### 🎯 Detection Results")
            output_image = gr.Image(
                label="Annotated Image",
                type="numpy",
                height=400
            )
            
            metrics_output = gr.HTML(label="Detailed Metrics")
    
    # Additional outputs
    with gr.Row():
        summary_output = gr.Textbox(
            label="📝 Text Summary",
            lines=10,
            max_lines=20
        )
    
    # Examples section
    gr.Markdown("### 💡 Example Images")
    gr.Examples(
        examples=[
            ["examples/pothole1.jpg", 0.5],
            ["examples/pothole2.jpg", 0.6],
            ["examples/pothole3.jpg", 0.4],
        ],
        inputs=[input_image, confidence_slider],
        outputs=[output_image, metrics_output, summary_output],
        fn=detect_potholes,
        cache_examples=False
    ) if Path("examples").exists() else None
    
    # Footer
    gr.HTML("""
        <div style="text-align: center; margin-top: 30px; padding: 20px; background: #f5f5f5; border-radius: 10px;">
            <p style="color: #666;">
                <strong>Powered by YOLOv11n-seg with Boundary-Aware Multi-Scale Training</strong><br>
                Built with ❤️ using Gradio and Ultralytics
            </p>
        </div>
    """)
    
    # Connect button to function
    detect_button.click(
        fn=detect_potholes,
        inputs=[input_image, confidence_slider],
        outputs=[output_image, metrics_output, summary_output]
    )

# ═══════════════════════════════════════════════════════════════════
# LAUNCH APPLICATION
# ═══════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    print("\n" + "="*80)
    print("🚀 LAUNCHING POTHOLE DETECTION APPLICATION")
    print("="*80)
    print(f"   Model: {MODEL_PATH}")
    print(f"   Device: {'GPU' if torch.cuda.is_available() else 'CPU'}")
    print("="*80 + "\n")
    
    app.launch(
        share=True,  # Creates public link
        server_name="0.0.0.0",  # Makes it accessible on network
        server_port=7860,
        show_error=True,
        # Enable queue for better handling of multiple users
        # max_threads=4
    )
# ```

# ---

# ## 📦 **Deployment Instructions**

# ### **File Structure:**
# ```
# pothole-detection-app/
# ├── app.py                    # Main Gradio application
# ├── requirements.txt          # Dependencies
# ├── best.pt                   # Your trained model checkpoint
# ├── README.md                 # Documentation
# └── examples/                 # (Optional) Example images
#     ├── pothole1.jpg
#     ├── pothole2.jpg
#     └── pothole3.jpg