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

# Google Drive file ID for best.pt
FILE_ID = "1VE4I-_OMoC-wzzo-udzhJ896-WdFwJjU"
FILE_URL = f"https://drive.google.com/uc?id={FILE_ID}"
MODEL_PATH = "best.pt"

# Download the model if not already present
if not os.path.exists(MODEL_PATH):
    st.info("Downloading model file from Google Drive...")
    try:
        response = requests.get(FILE_URL, stream=True)
        response.raise_for_status()
        with open(MODEL_PATH, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
        st.success("Model file downloaded successfully!")
    except Exception as e:
        st.error(f"Error downloading model: {e}")

# Load the YOLO model
try:
    model = YOLO(MODEL_PATH)
except Exception as e:
    st.error(f"Error loading model: {e}")

# Streamlit app
st.title("Roofy v5 basic test")
st.write("Upload an image and let the model detect objects.")

# Sliders for Confidence and Overlap thresholds
confidence_threshold = st.slider(
    "Confidence Threshold",
    min_value=0.0,
    max_value=1.0,
    value=0.25,
    step=0.01,
    help="Set the minimum confidence score for detections.",
)
overlap_threshold = st.slider(
    "Overlap Threshold",
    min_value=0.0,
    max_value=1.0,
    value=0.45,
    step=0.01,
    help="Set the maximum allowable overlap for non-max suppression.",
)

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_container_width=True)

    # Perform prediction
    with st.spinner("Processing..."):
        try:
            results = model.predict(
                np.array(image),
                conf=confidence_threshold,  # Set confidence threshold
                iou=overlap_threshold,  # Set overlap (IoU) threshold
            )
            st.write("Detection Results:")
            st.image(results[0].plot(), caption="Detections", use_container_width=True)
        except Exception as e:
            st.error(f"Error during prediction: {e}")