import gradio as gr
import supervision as sv
from ultralytics import YOLO
import numpy as np
import cv2  # for video to image conversion


model = YOLO('yolov8s.pt')
byte_tracker = sv.ByteTrack()
annotator = sv.BoxAnnotator()


def process_video(frame):
    results = model(frame)[0]
    detections = sv.Detections.from_ultralytics(results)
    detections = byte_tracker.update_with_detections(detections)
    labels = [
         f"#{tracker_id} {model.model.names[class_id]} {confidence:0.2f}"
         for _, _, confidence, class_id, tracker_id
         in detections
    ]
    yield  annotator.annotate(scene=frame.copy(),
                               detections=detections, labels=labels)


title = "Object Tracking (w/ YOLOv8)"
with gr.Blocks() as io:
    gr.Markdown(f"<center><h1>{title}</h1></center>")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(source='webcam', streaming=True)
            input_button = gr.Button()
        with gr.Column():
            output_image = gr.Image()
        input_image.change(process_video, inputs=[input_image], outputs=[output_image], show_progress=False)
        
io.queue()

io.launch()