import warnings
warnings.filterwarnings("ignore")

import streamlit as st
from ultralytics import YOLO
from PIL import Image
import tempfile
import os

# Load YOLO model
model = YOLO('yolov8n.pt')  # Make sure this is uploaded to the repo or use Hugging Face Hub path

st.title("🧠 YOLO Object Detection with Streamlit")
st.write("Upload an image to run real-time object detection.")

# Upload image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Convert uploaded file to PIL image
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Image", use_container_width=True)

    # Save to a temporary file for YOLO inference
    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_file:
        temp_path = tmp_file.name
        image.save(temp_path)

    # Confidence threshold slider
    conf = st.slider("Confidence Threshold", 0.0, 1.0, 0.25)

    # Run YOLO inference
    results = model(temp_path, conf=conf)

    # Display detection results
    st.image(results[0].plot(), caption="Detected Objects", use_container_width=True)

    # Detection details
    with st.expander("Detection Details"):
        for box in results[0].boxes:
            cls = model.names[int(box.cls)]
            conf_score = float(box.conf)
            st.write(f"**{cls}** — Confidence: {conf_score:.2f}")

    # Clean up temp file
    os.remove(temp_path)