import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO

# Load the model
model = YOLO('best.pt')

def get_intelligence(frame, box, label):
    x1, y1, x2, y2 = map(int, box.xyxy[0])
    w_px = x2 - x1
    h_px = y2 - y1
    
    # Distance estimation logic
    dist_m = (2.5 * 700) / w_px
    dist_cat = "Near" if dist_m < 15 else "Medium" if dist_m < 30 else "Far"
    
    # Aspect ratio logic for direction
    aspect_ratio = w_px / h_px
    direction = "Front/Rear" if aspect_ratio < 1.2 else "Side View"
    
    return dist_cat, direction

def predict(img):
    results = model(img, conf=0.75)
    annotated_img = img.copy()
    
    status_reports = []

    for r in results:
        for box in r.boxes:
            cls_id = int(box.cls[0])
            label = r.names[cls_id]
            
            dist, direction = get_intelligence(img, box, label)
            
            # Create a label for the report
            report = f"Detected {label} ({direction}) - Distance: {dist}"
            status_reports.append(report)
            
            # Draw on image
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            cv2.rectangle(annotated_img, (x1, y1), (x2, y2), (0, 255, 0), 3)
            cv2.putText(annotated_img, f"{label} {dist}", (x1, y1 - 10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
            
    return annotated_img, "\n".join(status_reports) if status_reports else "No emergency vehicles detected."

# Build the Gradio UI
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs=[gr.Image(type="numpy"), gr.Textbox(label="Intelligence Report")],
    title="EVobj: Emergency Vehicle Detection System",
    description="Upload an image of traffic to detect emergency vehicles and estimate proximity."
)

if __name__ == "__main__":
    demo.launch()