from ultralytics import YOLO
import gradio as gr
from gradio_webrtc import WebRTC

# Load the YOLO model
model = YOLO("yolo11n.pt")


def detect_objects(frame, conf_threshold=0.3):
    results = model(frame, conf=conf_threshold)
    annotated_frame = results[0].plot()
    return annotated_frame


css = """.my-group {max-width: 600px !important; max-height: 600px !important;}
.my-column {display: flex !important; justify-content: center !important; align-items: center !important;}"""

with gr.Blocks(css=css) as demo:
    gr.HTML(
        "<h1 style='text-align: center'>YOLOv10 Webcam Stream (Powered by WebRTC ⚡️)</h1>"
    )

    with gr.Column(elem_classes=["my-column"]):
        image = WebRTC(label="Stream")
        conf_threshold = gr.Slider(
            label="Confidence Threshold",
            minimum=0.0,
            maximum=1.0,
            step=0.05,
            value=0.30,
        )

        # Stream processing function
        image.stream(
            fn=detect_objects,
            inputs=[image, conf_threshold],
            outputs=[image],
            time_limit=10,
        )

demo.launch()