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

# ---------------- CONFIG ----------------
st.set_page_config(page_title="Pothole Detection", layout="wide")
st.title("🕳️ Pothole Detection using YOLO")
st.write("Upload an image — the model will detect potholes and mark them.")

# -------- Load YOLO Model --------------
@st.cache_resource
def load_model():
    try:
        model = YOLO("best.pt")  # your model file
        return model
    except Exception as e:
        st.error(f"Failed to load model: {e}")
        return None

model = load_model()

if model is None:
    st.stop()

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

if uploaded_file:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Image", use_container_width=True)

    with st.spinner("Detecting potholes... ⏳"):
        # Save temp file
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp:
            image.save(tmp.name)
            results = model(tmp.name)

        # Render result image
        result_img = results[0].plot()  # numpy array (BGR)

        # Convert BGR to RGB
        result_img_rgb = Image.fromarray(result_img[..., ::-1])

        st.image(result_img_rgb, caption="Detected Potholes ✅", use_container_width=True)

        # Download button
        result_path = "output_pothole.jpg"
        result_img_rgb.save(result_path)

        with open(result_path, "rb") as f:
            st.download_button("📥 Download Result", f, file_name="pothole_detected.jpg")