import streamlit as st
from PIL import Image
from pathlib import Path
from ultralytics import YOLO

# Load YOLOv5 model
model = YOLO('HandSignDetector.pt')  # Replace with the path to your best.pt model

# Set up Streamlit
st.title("YOLOv5 Object Detection with Streamlit")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

from PIL import Image, ImageDraw, ImageFont

def draw_bounding_box(image, box, class_label, probability):
    # Convert to ImageDraw format
    draw = ImageDraw.Draw(image)

    # Draw bounding box
    draw.rectangle(box, outline="red", width=3)

    # Add class label and probability
    label = f"{class_label}: {probability:.2f}"
    font = ImageFont.load_default()
    text_width, text_height = draw.textsize(label, font)

    # Calculate position to center the text inside the bounding box
    text_position = ((box[0] + box[2]) - text_width) / 2, box[3] + 5

    # Draw text on the image
    draw.text(text_position, label, font=font, fill="red")
    return draw
    
if uploaded_file is not None:
    # Read the uploaded image
    image = Image.open(uploaded_file)

    # Inference
    results = model(image)

    # Display the image with bounding boxes
    st.image(image, channels="RGB", caption="Object Detection Result", use_column_width=True)

    # Display probability and class for each box
    for det in results:
        box = det.boxes.xyxy
        cls = det.boxes.cls
        score = det.boxes.conf

        # Draw box on the image
        st.image(draw_bounding_box(image,box,cls,score))