import streamlit as st
from PIL import Image
import torch
import os

# Load the trained YOLOv5 model
model = torch.hub.load("ultralytics/yolov5", "custom", path="best.pt")  # Load custom model

# Define the prediction function
def predict(image):
    # Perform inference on the uploaded image
    results = model(image)  # Runs YOLOv5 model on the uploaded image
    results_img = results.render()[0]  # Get image with bounding boxes drawn
    return Image.fromarray(results_img)

# Get example images from the images folder
def get_example_images():
    examples = []
    image_folder = "images"
    for filename in os.listdir(image_folder):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
            examples.append(os.path.join(image_folder, filename))
    return examples

# Streamlit UI for Seatbelt Detection with YOLOv5
st.title("Seatbelt Detection with YOLO")
st.markdown("Upload an image to detect Seat-Belt-Detection.")

# Allow the user to upload an image
uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_image is not None:
    # Open the uploaded image using PIL
    image = Image.open(uploaded_image)

    # Display the uploaded image
    st.image(image, caption="Uploaded Image", use_container_width=True)

    # Run the model prediction
    st.subheader("Prediction Results:")
    result_image = predict(image)
    
    # Display the result image with bounding boxes
    st.image(result_image, caption="Detected Image", use_container_width=True)

# Optionally, show example images from the folder
if st.checkbox('Show example images'):
    example_images = get_example_images()
    for example_image in example_images:
        img = Image.open(example_image)
        st.image(img, caption=os.path.basename(example_image), use_container_width=True)