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

# --- Model Loading ---
model = None
try:
    # Load the trained YOLOv5 model
    # Suppress verbose output from torch.hub.load
    with st.spinner("Loading YOLOv5 model..."):
        model = torch.hub.load("ultralytics/yolov5", "custom", path="best.pt", force_reload=False, verbose=False)
    st.success("Model loaded successfully!")
except Exception as e:
    st.error(f"Error loading the YOLOv5 model: {e}")
    st.warning("Please ensure 'best.pt' is in the correct directory and accessible, and that you have an internet connection to download YOLOv5 dependencies.")

# Define the prediction function
def predict(image):
    if model is None:
        st.error("Model is not loaded. Cannot perform prediction.")
        return None
    try:
        # 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)
    except Exception as e:
        st.error(f"Error during prediction: {e}")
        return None

# Get example images from the images folder
def get_example_images():
    examples = []
    image_folder = "images"
    if not os.path.exists(image_folder):
        st.warning(f"Example image folder '{image_folder}' not found.")
        return []
    try:
        for filename in os.listdir(image_folder):
            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                examples.append(os.path.join(image_folder, filename))
        return examples
    except Exception as e:
        st.error(f"Error loading example images: {e}")
        return []

# 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:
    image = None
    try:
        # 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 if prediction was successful
        if result_image is not None:
            st.image(result_image, caption="Detected Image", use_container_width=True)
    except Exception as e:
        st.error(f"Error processing uploaded image: {e}")
        st.warning("Please ensure the uploaded file is a valid image.")

# Optionally, show example images from the folder
if st.checkbox('Show example images'):
    example_images = get_example_images()
    if example_images:
        for example_image in example_images:
            try:
                img = Image.open(example_image)
                st.image(img, caption=os.path.basename(example_image), use_container_width=True)
            except Exception as e:
                st.error(f"Could not load example image {os.path.basename(example_image)}: {e}")
    else:
        st.info("No example images found or an error occurred while loading them.")