app.py

"""
YOLOv8 Traffic Bird's Eye View Detection Application

This application provides a web-based interface for traffic object detection using 
a specialized YOLOv8 model trained on aerial/bird's eye view imagery.
Users can upload aerial images or select from examples to receive real-time 
traffic detection results with visual annotations and structured data output.
"""

import logging
import os
import sys
import tempfile
from datetime import datetime
from typing import Tuple, Optional

import cv2
import gradio as gr
import numpy as np
import pandas as pd
from PIL import Image
from ultralytics import YOLO

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global model instance
model: Optional[YOLO] = None


def initialize_model() -> bool:
    """
    Initialize YOLOv8 model with error handling and validation.

    Returns:
        bool: True if model loaded successfully, False otherwise
    """
    global model
    try:
        logger.info("Loading traffic bird's eye view YOLOv8 model...")
        model = YOLO("./traffic_birdseye.pt")
        logger.info("Traffic bird's eye view YOLOv8 model loaded successfully")
        return True
    except Exception as e:
        logger.error("Failed to load traffic bird's eye view YOLOv8 model: %s", str(e))
        model = None
        return False


def preprocess_image(image: Image.Image) -> np.ndarray:
    """
    Convert PIL Image to NumPy array for YOLOv8 processing.

    Args:
        image (Image.Image): PIL Image object

    Returns:
        np.ndarray: NumPy array in RGB format
    """
    try:
        # Convert PIL Image to RGB if not already
        if image.mode != 'RGB':
            image = image.convert('RGB')

        # Convert to NumPy array
        image_array = np.array(image)

        logger.info("Image preprocessed: shape %s", image_array.shape)
        return image_array

    except Exception as e:
        logger.error("Error preprocessing image: %s", str(e))
        raise


def detect_objects(image: Image.Image, conf_threshold: float = 0.25) -> Tuple[Image.Image, pd.DataFrame, str]:
    """
    Main object detection function that processes images through YOLOv8.

    Args:
        image (Image.Image): PIL Image object to process
        conf_threshold (float): Confidence threshold for filtering detections (0.0 to 1.0)

    Returns:
        Tuple[Image.Image, pd.DataFrame, str]: Annotated image, detection results table, and CSV file path
    """
    if model is None:
        logger.error("YOLOv8 model not loaded")
        # Return original image and empty dataframe
        empty_df = pd.DataFrame(columns=['class', 'confidence', 'x1', 'y1', 'x2', 'y2'])
        return image, empty_df, None

    try:
        # Preprocess the image
        logger.info("Starting object detection with confidence threshold: %.2f", conf_threshold)
        processed_image = preprocess_image(image)

        # Run YOLOv8 inference with confidence threshold
        results = model(processed_image, conf=conf_threshold)

        # Parse detection results
        annotated_image, detection_data = parse_detection_results(results[0], image, conf_threshold)

        # Create CSV file for download
        csv_file_path = create_csv_download(detection_data, conf_threshold)

        logger.info("Object detection completed successfully with %d detections", len(detection_data))
        return annotated_image, detection_data, csv_file_path

    except Exception as e:
        logger.error("Error during object detection: %s", str(e))
        # Return original image and empty dataframe on error
        empty_df = pd.DataFrame(columns=['class', 'confidence', 'x1', 'y1', 'x2', 'y2'])
        return image, empty_df, None


def parse_detection_results(results, original_image: Image.Image, conf_threshold: float = 0.25) -> Tuple[Image.Image, pd.DataFrame]:
    """
    Parse YOLOv8 detection results to extract classes, confidence scores, and bounding boxes.
    
    Args:
        results: YOLOv8 detection results object
        original_image (Image.Image): Original PIL image for annotation
        conf_threshold (float): Confidence threshold used for filtering
        
    Returns:
        Tuple[Image.Image, pd.DataFrame]: Annotated image and structured detection data
    """
    try:
        # Extract detection data
        detection_list = []
        
        # Check if any detections were found
        if results.boxes is not None and len(results.boxes) > 0:
            # Get class names from the model
            class_names = results.names
            
            # Extract boxes, classes, and confidence scores
            boxes = results.boxes.xyxy.cpu().numpy()  # x1, y1, x2, y2 format
            classes = results.boxes.cls.cpu().numpy()
            confidences = results.boxes.conf.cpu().numpy()
            
            # Process each detection (already filtered by YOLOv8 based on conf_threshold)
            for i in range(len(boxes)):
                class_id = int(classes[i])
                class_name = class_names[class_id]
                confidence = float(confidences[i])
                x1, y1, x2, y2 = boxes[i]
                
                # Create detection record with proper rounding
                detection_record = {
                    'class': class_name,
                    'confidence': round(confidence, 3),  # Round to 3 decimal places
                    'x1': round(float(x1), 1),          # Round to 1 decimal place
                    'y1': round(float(y1), 1),
                    'x2': round(float(x2), 1),
                    'y2': round(float(y2), 1)
                }
                
                detection_list.append(detection_record)
                
            logger.info("Parsed %d detections above confidence threshold %.2f", len(detection_list), conf_threshold)
        else:
            logger.info("No objects detected above confidence threshold %.2f", conf_threshold)
        
        # Create DataFrame from detection list with proper columns
        if detection_list:
            detection_df = pd.DataFrame(detection_list)
        else:
            # Create empty DataFrame with correct columns
            detection_df = pd.DataFrame(columns=['class', 'confidence', 'x1', 'y1', 'x2', 'y2'])
        
        # Generate annotated image using YOLOv8's built-in plotting
        annotated_image = generate_annotated_image(results, original_image)
        
        return annotated_image, detection_df
        
    except Exception as e:
        logger.error("Error parsing detection results: %s", str(e))
        # Return original image and empty dataframe on error
        empty_df = pd.DataFrame(columns=['class', 'confidence', 'x1', 'y1', 'x2', 'y2'])
        return original_image, empty_df


def generate_annotated_image(results, original_image: Image.Image) -> Image.Image:
    """
    Generate annotated image with bounding boxes and labels using YOLOv8's plot method.
    
    Args:
        results: YOLOv8 detection results object
        original_image (Image.Image): Original PIL image
        
    Returns:
        Image.Image: Annotated PIL image with bounding boxes and labels in RGB format
    """
    try:
        # Use YOLOv8's built-in plot method to generate annotations
        # This automatically adds bounding boxes, class labels, and confidence scores
        annotated_array = results.plot(
            conf=True,  # Show confidence scores
            labels=True,  # Show class labels
            boxes=True,  # Show bounding boxes
            line_width=2,  # Line width for bounding boxes
            font_size=12  # Font size for labels
        )
        
        # CRITICAL: YOLOv8's plot() method ALWAYS returns BGR format, but Gradio expects RGB
        # Force BGR to RGB conversion regardless of array shape
        #logger.info("Converting BGR to RGB for proper color display")
        #annotated_rgb = cv2.cvtColor(annotated_array, cv2.COLOR_BGR2RGB)
        
        # Ensure the array is in the correct data type
        annotated_array = annotated_array.astype(np.uint8)
        
        # Convert NumPy array back to PIL Image
        annotated_image = Image.fromarray(annotated_array)
        
        # Verify the conversion worked
        logger.info("Generated annotated image successfully with dimensions: %s, mode: %s", 
                   annotated_image.size, annotated_image.mode)
        return annotated_image
        
    except Exception as e:
        logger.error("Error generating annotated image: %s", str(e))
        # Return original image on error to maintain aspect ratio
        return original_image


def create_csv_download(detection_data: pd.DataFrame, conf_threshold: float = 0.25) -> str:
    """
    Create a CSV file from detection results for download.
    
    Args:
        detection_data (pd.DataFrame): Detection results dataframe
        conf_threshold (float): Confidence threshold used for filtering
        
    Returns:
        str: Path to the created CSV file
    """
    try:
        # Create a temporary CSV file
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        csv_filename = f"detection_results_{timestamp}.csv"
        
        # Use a temporary directory
        temp_dir = tempfile.gettempdir()
        csv_path = os.path.join(temp_dir, csv_filename)
        
        # Save DataFrame to CSV with additional metadata
        if not detection_data.empty:
            # Add metadata to the CSV
            with open(csv_path, 'w', newline='', encoding='utf-8') as f:
                f.write(f"# YOLOv8 Object Detection Results\n")
                f.write(f"# Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"# Confidence threshold: {conf_threshold:.2f} ({conf_threshold*100:.0f}%)\n")
                f.write(f"# Total detections: {len(detection_data)}\n")
                f.write("# Columns: class, confidence, x1, y1, x2, y2\n")
                f.write("# Coordinates are in pixels (x1,y1 = top-left, x2,y2 = bottom-right)\n")
                f.write("#\n")
            
            # Append the actual data
            detection_data.to_csv(csv_path, mode='a', index=False)
        else:
            # Create empty CSV with headers
            detection_data.to_csv(csv_path, index=False)
        
        logger.info("CSV file created successfully: %s", csv_path)
        return csv_path
        
    except Exception as e:
        logger.error("Error creating CSV file: %s", str(e))
        return None


def create_gradio_interface():
    """
    Create Gradio Interface with image upload, confidence slider, and triple output.
    
    Returns:
        gr.Interface: Configured Gradio interface
    """
    # Create the Gradio interface
    demo = gr.Interface(
        fn=detect_objects,
        inputs=[
            gr.Image(
                type="pil",
                label="Upload Image for Object Detection",
                sources=["upload", "clipboard"],
                height=400
            ),
            gr.Slider(
                minimum=0.01,
                maximum=1.0,
                value=0.25,
                step=0.01,
                label="Confidence Threshold",
                info="Only show detections above this confidence level (0.01 = 1%, 1.0 = 100%)"
            )
        ],
        outputs=[
            gr.Image(
                type="pil",
                label="Annotated Image with Detections",
                height=400,
                show_label=True
            ),
            gr.Dataframe(
                label="Detection Results",
                headers=["Class", "Confidence", "X1", "Y1", "X2", "Y2"],
                datatype=["str", "number", "number", "number", "number", "number"],
                col_count=(6, "fixed"),
                row_count=(10, "dynamic"),
                interactive=False
            ),
            gr.File(
                label="Download Results as CSV",
                visible=True
            )
        ],
        title="🚗 Traffic Bird's Eye View Object Detection",
        description="""
        **Upload an aerial/bird's eye view image to detect traffic objects!**
        
        This application uses a specialized YOLOv8 model trained for traffic detection from aerial perspectives.
        Perfect for analyzing drone footage, satellite images, or overhead traffic cameras.
        
        **Supported formats:** JPG, PNG, JPEG
        
        **Features:** 
        - 🚗 Specialized traffic detection model
        - 🎚️ Adjustable confidence threshold slider
        - 📥 Download detection results as CSV
        - 📊 Real-time filtering based on confidence scores
        - 🖼️ Example images to try out
        """,
        article="""
        ### How it works:
        1. **Upload** an aerial/bird's eye view image or **select from examples** below
        2. **Adjust** the confidence threshold slider to filter detections (default: 0.25 = 25%)
        3. **View** the annotated image with bounding boxes and labels for detected traffic objects
        4. **Analyze** the detection data in the table with confidence scores and coordinates
        5. **Download** the results as a CSV file for further analysis
        
        **Best Results:** This model works best with:
        - 🚁 Drone footage of roads and traffic
        - 🛰️ Aerial photography of urban areas
        - 📹 Overhead traffic camera feeds
        - 🗺️ Bird's eye view street scenes
        
        **Confidence Threshold Guide:**
        - **0.01-0.20**: Very permissive - shows many detections, including uncertain ones
        - **0.25-0.50**: Balanced - good mix of accuracy and detection count (recommended)
        - **0.50-0.80**: Conservative - only high-confidence detections
        - **0.80-1.00**: Very strict - only extremely confident detections
        
        Try the example images below to see the model in action!
        """,
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1200px !important;
        }
        .output-image {
            height: 400px !important;
        }
        .slider-container {
            margin: 10px 0;
        }
        """,
        examples=[
            ["examples/13sbt-aerial-shot-of-white-convertible-car-riding-throug-F5ZX2CK-12.jpg", 0.25],
            ["examples/7sbt-cars-on-the-road-aerial-view-QLX4X43-1481.jpg", 0.25],
            ["examples/13sbt-car-traffic-yellow-autumn-forest-nature-road-lands-SFXKS2M-86.jpg", 0.25],
            ["examples/7sbt-cars-ride-on-the-road-slow-motion-kyiv-ukraine-aer-W6BZ8YJ-5520.jpg", 0.25],
            ["examples/13sbt-DJI_0123-3.jpg", 0.25],
            ["examples/2022teknofest-DJI_0309-86.jpg", 0.25]
        ],
        cache_examples=False,
        flagging_mode="never"
    )
    
    return demo


if __name__ == "__main__":
    # Initialize model on startup
    if not initialize_model():
        logger.error("Failed to initialize YOLOv8 model. Application cannot start.")
        sys.exit(1)
    
    logger.info("YOLOv8 model initialized successfully")
    
    # Create and launch Gradio interface
    try:
        demo = create_gradio_interface()
        logger.info("Gradio interface created successfully")
        
        # Launch the application
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False,
            debug=False,
            show_error=True,
            quiet=False
        )
        
    except Exception as e:
        logger.error("Error launching Gradio interface: %s", str(e))
        sys.exit(1)