app.py

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 YOLO model to detect benthic megafauna in images.")

# 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**
    """
)