import streamlit as st
from ultralytics import YOLO
from PIL import Image
import numpy as np
import os

# Load the model
model_path = "./models/best.pt"
model = YOLO(model_path)

# Streamlit app
st.title("YOLOv11 Object Detection")
st.write("Upload an image and let the model detect objects.")

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
if uploaded_file:
    # Read and display the image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Perform prediction
    with st.spinner("Processing..."):
        results = model.predict(np.array(image))
    
    # Display results
    st.write("Detection Results:")
    st.image(results[0].plot(), caption="Detections", use_column_width=True)