import gradio as gr
import cv2
from ultralytics import YOLO

# Load the pretrained YOLOv5 model
yolo_model = YOLO("yolov5s.pt")  # Load the small version of YOLOv5

# Function to perform object detection with a configurable confidence threshold
def detect_objects(frame, confidence_threshold=0.5):
    # Convert the frame from BGR (OpenCV) to RGB
    image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    
    # Perform inference with YOLO, passing the confidence threshold
    results = yolo_model(image, conf=confidence_threshold)  # Set confidence threshold
    
    # Draw bounding boxes and labels on the image
    annotated_image = results.plot()  # This automatically draws boxes and labels
    
    # Convert the image back to BGR for displaying in Gradio
    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR)
    
    return annotated_image

# Gradio interface to use the webcam for real-time object detection
# Added a slider for the confidence threshold
iface = gr.Interface(fn=detect_objects, 
                     inputs=[
                         gr.Video(source="webcam", type="numpy"),  # Webcam input
                         gr.Slider(minimum=0.0, maximum=1.0, default=0.5, label="Confidence Threshold")  # Confidence slider
                     ],  
                     outputs="image")  # Show output image with bounding boxes

iface.launch()