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

# Konfigurasi Model
MODEL_PATH = "best.pt"
CLASS_NAMES = ["bag", "person-static-object"]

# Contoh gambar lokal dalam folder yang sama
SAMPLE_IMAGES = {
    "Contoh 1": "sample1.jpg",
    "Contoh 2": "sample2.jpg"
}

# Muat Model
@st.cache_resource
def load_model():
    return YOLO(MODEL_PATH)

# Streamlit UI
st.title("Deteksi objek tertinggal 🚀")
st.write("!! Projek ini merupakan simulasi deteksi objek yang tertinggal, untuk saat ini hanya bisa mendeteksi dari sebuah gambar dikarenakan keterbatasan sumber daya !!")

# Pilihan gambar
option = st.radio("Pilih sumber gambar:", ["Upload Gambar", "Pilih Contoh Gambar"])

if option == "Upload Gambar":
    uploaded_file = st.file_uploader("Upload gambar Anda...", type=["jpg", "png", "jpeg"])
    image = Image.open(uploaded_file).convert("RGB") if uploaded_file else None
else:
    selected_sample = st.selectbox("Pilih contoh gambar:", list(SAMPLE_IMAGES.keys()))
    try:
        image_path = SAMPLE_IMAGES[selected_sample]
        image = Image.open(image_path).convert("RGB")
    except Exception as e:
        st.error(f"Gagal memuat gambar: {str(e)}")
        image = None

if image:
    st.image(image, caption="Gambar Input", use_container_width=True)
    
    if st.button("Deteksi Objek"):
        with st.spinner("Memproses..."):
            try:
                model = load_model()
                results = model.predict(image)
                
                # Visualisasi hasil
                res_plotted = results[0].plot()[:, :, ::-1]
                st.image(res_plotted, caption="Hasil Deteksi", use_container_width=True)
                
                # Tampilkan statistik
                boxes = results[0].boxes
                st.success(f"✅ Objek Terdeteksi: {len(boxes)}")
                
                # Tampilkan detail
                if len(boxes) > 0:
                    st.subheader("Detail Deteksi:")
                    for i, box in enumerate(boxes):
                        cls = CLASS_NAMES[int(box.cls)]
                        conf = box.conf[0].item()
                        st.write(f"{i+1}. {cls} (confidence: {conf:.2f})")
                        
            except Exception as e:
                st.error(f"❌ Error: {str(e)}")