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

st.set_page_config(page_title="OceanCV Object Detection", page_icon="🦑")

st.title("OceanCV FirstPass Detector")
st.write("Upload marine imagery to detect underwater objects.")

@st.cache_resource
def load_model():
    return YOLO("OceanCV_FirstPass.pt")

model = load_model()

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

conf = st.slider("Confidence Threshold", 0.1, 1.0, 0.5)
iou = st.slider("IOU Threshold", 0.1, 1.0, 0.45)

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    
    if st.button("Detect Objects"):
        with st.spinner("Analyzing..."):
            results = model.predict(image, conf=conf, iou=iou, imgsz=1024)
            
            res_plotted = results[0].plot()
            res_rgb = cv2.cvtColor(res_plotted, cv2.COLOR_BGR2RGB)
            
            st.image(res_rgb, caption="Processed Image", use_column_width=True)
            
            st.write(f"Detected {len(results[0].boxes)} objects.")