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

st.title("Object Detection (YOLOv8)")

@st.cache_resource
def load_model():
    return YOLO("models/best.pt")

model = load_model()

confidence = st.slider("Confidence Threshold", 0.1, 1.0, 0.25)

uploaded_file = st.file_uploader("Upload Image", type=["jpg","jpeg","png"])

if uploaded_file:
    try:
        image = Image.open(uploaded_file)
        img_array = np.array(image)

        results = model.predict(img_array, conf=confidence)
        annotated = results[0].plot()

        st.image(annotated, width="stretch")

    except Exception as e:
        st.error(f"Detection failed: {e}")