# -*- coding: utf-8 -*-
import os
import cv2
import time
import numpy as np
import gradio as gr
from PIL import Image
from ultralytics import YOLO

model = YOLO("yolov8n.pt")

def detect(image, conf_threshold):
    if image is None:
        return None, "No image provided"

    temp_path = "temp_input.jpg"
    Image.fromarray(image).save(temp_path)

    start = time.time()
    results = model(temp_path, conf=conf_threshold, verbose=False)
    latency = round((time.time() - start) * 1000, 2)

    annotated = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)

    boxes = results[0].boxes
    class_counts = {}
    for box in boxes:
        cls_name = model.names[int(box.cls[0])]
        class_counts[cls_name] = class_counts.get(cls_name, 0) + 1

    if not class_counts:
        summary = f"Latency: {latency}ms | No objects detected above {int(conf_threshold*100)}% confidence"
    else:
        summary = [f"Latency: {latency}ms | Total: {len(boxes)} detections"]
        for cls, count in class_counts.items():
            summary.append(f"  - {cls}: {count}")
        summary = "\n".join(summary)

    return annotated, summary

with gr.Blocks(theme=gr.themes.Soft(), title="YOLOv8 Object Detection") as app:
    gr.Markdown("# YOLOv8 Real-Time Object Detection")
    gr.Markdown("Detect objects in images or via live webcam using YOLOv8n.")

    with gr.Tab("Image Upload"):
        with gr.Row():
            img_input  = gr.Image(type="numpy", label="Upload Image")
            img_output = gr.Image(type="numpy", label="Detection Results")
        img_conf    = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Confidence Threshold")
        img_summary = gr.Textbox(label="Detection Summary")
        img_btn     = gr.Button("Detect Objects")
        img_btn.click(fn=detect, inputs=[img_input, img_conf], outputs=[img_output, img_summary])

    with gr.Tab("Webcam Live"):
        gr.Markdown("Allow camera access. Detections appear automatically every 0.5 seconds.")
        with gr.Row():
            cam_input  = gr.Image(sources=["webcam"], type="numpy", label="Live Feed", streaming=True)
            cam_output = gr.Image(type="numpy", label="Detection Results")
        cam_conf    = gr.Slider(0.1, 0.9, value=0.25, step=0.05, label="Confidence Threshold")
        cam_summary = gr.Textbox(label="Detection Summary")
        cam_input.stream(
            fn=detect,
            inputs=[cam_input, cam_conf],
            outputs=[cam_output, cam_summary],
            time_limit=60,
            stream_every=0.5
        )

app.launch(server_name="0.0.0.0", server_port=7860, share=False)