Update app.py

app.py

@@ -1,7 +1,9 @@
+import streamlit as st
 import glob
-import gradio as gr
 from ultralytics import YOLO
+from PIL import Image
 
+# Load the YOLO model
 model_path = "best.pt"
 model = YOLO(model_path)
 
@@ -9,22 +11,39 @@ PREDICT_KWARGS = {
     "conf": 0.15,
 }
 
-def run(image_path):
-    results = model.predict(image_path, **PREDICT_KWARGS)
-    return results[0].plot()[:, :, ::-1]  # reverse channels for gradio
-
-title = "OOI RCA Digital Still Camera Benthic Megafauna Detector"
-description = ""
+# 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")
 
-interface = gr.Interface(
-    run,
-    inputs=[gr.components.Image(type="filepath")],
-    outputs=gr.components.Image(type="numpy"),
-    title=title,
-    description=description,
-    examples=examples,
-)
+# Display list of example images to choose from
+st.sidebar.title("Example Images")
+selected_example = st.sidebar.selectbox("Select an example image", examples)
+
+# Display the selected example image
+if selected_example:
+    st.image(selected_example, caption="Selected Example Image", use_column_width=True)
 
-interface.queue().launch()
+# 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)
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+    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 on the selected image
+if image_path is not None:
+    results = model.predict(image_path, **PREDICT_KWARGS)
+    st.image(results[0].plot()[:, :, ::-1], caption="Predicted Image", use_column_width=True)
+else:
+    st.write("Please upload an image or select an example to proceed.")