import torch
import numpy as np
import gradio as gr
import cv2
import time
import os
from pathlib import Path

# Create cache directory for models
os.makedirs("models", exist_ok=True)

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

# Load YOLOv5x model
model_path = Path("models/yolov5x.pt")
if model_path.exists():
    print(f"Loading model from cache: {model_path}")
    model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True, source="local", path=str(model_path)).to(device)
else:
    print("Downloading YOLOv5x model and caching...")
    model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True).to(device)
    torch.save(model.state_dict(), model_path)

# Model configurations
model.conf = 0.3  # Confidence threshold
model.iou = 0.3   # IoU threshold
model.classes = None  # Detect all classes

if device.type == "cuda":
    model.half()
else:
    torch.set_num_threads(os.cpu_count())

model.eval()

np.random.seed(42)
colors = np.random.uniform(0, 255, size=(len(model.names), 3))

total_inference_time = 0
inference_count = 0

def detect_objects(image):
    global total_inference_time, inference_count
    
    if image is None:
        return None
    
    start_time = time.time()
    output_image = image.copy()
    input_size = 640
    
    with torch.no_grad():
        results = model(image, size=input_size)
    
    inference_time = time.time() - start_time
    total_inference_time += inference_time
    inference_count += 1
    avg_inference_time = total_inference_time / inference_count
    
    detections = results.pred[0].cpu().numpy()
    
    for *xyxy, conf, cls in detections:
        x1, y1, x2, y2 = map(int, xyxy)
        class_id = int(cls)
        color = colors[class_id].tolist()
        
        # Thicker bounding boxes
        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
        
        label = f"{model.names[class_id]} {conf:.2f}"
        font_scale, font_thickness = 0.9, 2  # Increased for better readability
        (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
        
        cv2.rectangle(output_image, (x1, y1 - h - 10), (x1 + w + 10, y1), color, -1)
        cv2.putText(output_image, label, (x1 + 5, y1 - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness, lineType=cv2.LINE_AA)
    
    fps = 1 / inference_time
    
    # Stylish FPS display
    overlay = output_image.copy()
    cv2.rectangle(overlay, (10, 10), (300, 80), (0, 0, 0), -1)
    output_image = cv2.addWeighted(overlay, 0.6, output_image, 0.4, 0)
    cv2.putText(output_image, f"FPS: {fps:.2f}", (20, 40),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, lineType=cv2.LINE_AA)
    cv2.putText(output_image, f"Avg FPS: {1/avg_inference_time:.2f}", (20, 70),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, lineType=cv2.LINE_AA)
    
    return output_image

example_images = ["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
os.makedirs("examples", exist_ok=True)

with gr.Blocks(title="Optimized YOLOv5 Object Detection") as demo:
    gr.Markdown("""
    # Optimized YOLOv5 Object Detection
    Detects objects using YOLOv5 with enhanced visualization and FPS tracking.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="Input Image", type="numpy")
            submit_button = gr.Button("Submit", variant="primary")
            clear_button = gr.Button("Clear")
        
        with gr.Column(scale=1):
            output_image = gr.Image(label="Detected Objects", type="numpy")
    
    gr.Examples(
        examples=example_images,
        inputs=input_image,
        outputs=output_image,
        fn=detect_objects,
        cache_examples=True
    )
    
    submit_button.click(fn=detect_objects, inputs=input_image, outputs=output_image)
    clear_button.click(lambda: (None, None), None, [input_image, output_image])

demo.launch()