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

# Load the YOLO model
model_path = "best.pt"
model = YOLO(model_path)

PREDICT_KWARGS = {
    "conf": 0.15,
}

# Title and description for the app
st.title("OOI RCA Digital Still Camera Benthic Megafauna Detector")
st.write("This app uses a YOLOv11 model to detect benthic megafauna in images taken by the RCA OOI Digital Still Camera placed at Southern Hydrate Ridge to provide near real-time biological data. Please select an image from the sidebar on the left to try it out! This model was trained by University of Washington researchers on a RTX 3090 GPU for 100 epochs using MBARIs YOLOv8 model FathomNet/MBARI-315k-yolov8 as a starting point.The dataset used was first annotated by the afrorementioned MBARI model before being preprocessed both algorithmically and by hand by Dr. Katie Bigham and Atticus Carter. Classes are constrained to morphotypes as exact species are yet to be fully IDed and integrated. This models ideal image input size for training is 1024x1024. ")

# Load example images
examples = glob.glob("images/*.png")

# Sidebar for selecting example images
st.sidebar.title("Example Images")
selected_example = st.sidebar.selectbox("Select an example image", examples)

# File uploader for custom images
uploaded_file = st.file_uploader("Or upload an image", type=["png", "jpg", "jpeg"])

# Select which image to use for prediction
if uploaded_file is not None:
    image = Image.open(uploaded_file)
    image_path = uploaded_file
elif selected_example:
    image = Image.open(selected_example)
    image_path = selected_example
else:
    image = None
    image_path = None

# Run the YOLO model and display the results side by side
if image_path is not None:
    results = model.predict(image_path, **PREDICT_KWARGS)
    prediction_image = results[0].plot()[:, :, ::-1]  # Processed image with predictions

    # Create two columns for side-by-side display
    col1, col2 = st.columns(2)

    with col1:
        st.image(image, caption="Selected Image", use_column_width=True)
    
    with col2:
        st.image(prediction_image, caption="Predicted Image", use_column_width=True)
else:
    st.write("Please upload an image or select an example to proceed.")

# Footer with copyright and image credit
st.markdown(
    """
    ---
    **© 2024 Atticus Carter, [https://oceancv.org/](https://oceancv.org/)**  
    **Image Credit: Ocean Observatories Initiative 2022**
    """
)