Update app.py

app.py

@@ -1,281 +1,281 @@
-# Import required modules
-import streamlit as st
-from ultralytics import YOLO
-from PIL import Image
-import os
-import json
-import logging
-import tempfile
-import pandas as pd
-import matplotlib.pyplot as plt
-
-st.set_page_config(
-    page_title="Fish Detector",
-    page_icon="🐟",
-    layout="wide"
-)
-sample_images_folder = "./images/sample_images"
-logging.basicConfig(level=logging.INFO)
-
-# Model loading
-model_folder = "./models"
-st.sidebar.title("🐟 Fish or No Fish Detector")
-st.sidebar.markdown("""
-### For more information:
-- Contact: Michael.Akridge@NOAA.gov
-- Visit the [GitHub repository](https://github.com/MichaelAkridge-NOAA/Fish-or-No-Fish-Detector/)
-""")
-# Display model links
-st.sidebar.markdown("### Model Links")
-st.sidebar.markdown("- [YOLO11 Fish Detector - Grayscale](https://huggingface.co/akridge/yolo11-fish-detector-grayscale)")
-st.sidebar.markdown("- [YOLO11 Segment Fish - Grayscale](https://huggingface.co/akridge/yolo11-segment-fish-grayscale)")
-model_name = st.sidebar.selectbox("Select a YOLO model", os.listdir(model_folder))
-model_path = os.path.join(model_folder, model_name)
-if not os.path.exists(model_path):
-    st.error(f"Model file not found at {model_path}. Please check your setup.")
-    st.stop()
-model = YOLO(model_path)
-
-# Sidebar configuration
-st.sidebar.header("Model Parameters")
-confidence = st.sidebar.slider("Detection Confidence Threshold", 0.0, 1.0, 0.35)
-final_confidence = st.sidebar.slider("Final Yes/No Confidence Threshold", 0.0, 1.0, 0.5)
-
-# Title and description
-st.title("🐟 Fish or No Fish Detector")
-st.write("""
-Is there a fish 🐟 or not? Upload one or more images to detect fish. Using a trained [Ultralytics YOLO11 Model](https://github.com/ultralytics/ultralytics) for its object detection capabilities.
-
-""")
-
-# Custom CSS for button and uploader alignment
-st.markdown("""
-    <style>
-    .custom-file-uploader {
-        display: flex;
-        align-items: center;
-        margin-top: -10px; /* Adjust to move button closer */
-        justify-content: flex-start;
-    }
-    .css-1cpxqw2 {
-        flex-grow: 1;  /* Let file uploader take remaining space */
-    }
-    .sample-button {
-        font-size: 14px;
-        padding: 8px;
-        background-color: #007BFF;
-        color: white;
-        border: none;
-        border-radius: 5px;
-        cursor: pointer;
-        margin-left: 10px;
-        height: 38px; /* Ensure button matches uploader height */
-    }
-    .sample-button:hover {
-        background-color: #0056b3;
-    }
-    </style>
-""", unsafe_allow_html=True)
-
-# Custom CSS for default button styling
-st.markdown("""
-    <style>
-    .stButton>button, .stDownloadButton>button {
-        width: 100%;
-        padding: 10px;
-        border-radius: 5px;
-        font-size: 18px;
-        font-weight: bold;
-        background-color: #007BFF;
-        color: white;
-        border: none;
-        cursor: pointer;
-    }
-    .stButton>button:hover, .stDownloadButton>button:hover {
-        background-color: #0056b3;
-    }
-    </style>
-""", unsafe_allow_html=True)
-# Load sample images function
-def load_sample_images():
-    return [os.path.join(sample_images_folder, img) for img in os.listdir(sample_images_folder) if img.lower().endswith(('png', 'jpg', 'jpeg'))]
-
-# Prediction function
-def run(image_path):
-    results = model.predict(image_path, conf=confidence)
-    boxes = []
-    fish_count = 0
-    confidences = []
-
-    for box in results[0].boxes:
-        x1, y1, x2, y2 = box.xyxy[0].tolist()
-        conf = box.conf[0].item()
-        class_id = int(box.cls[0].item())
-        class_label = model.names[class_id].lower()  # Normalize to lowercase
-
-        if class_label == "fish" and conf > confidence:
-            fish_count += 1
-        confidences.append(conf)
-
-        boxes.append({"x1": x1, "y1": y1, "x2": x2, "y2": y2, "confidence": conf, "class_id": class_id, "class_label": class_label})
-
-    return results[0].plot()[:, :, ::-1], {"fish_count": fish_count, "confidences": confidences}
-
-# Process images function with directory creation
-# Reusable function to handle multiple image uploads and display results
-def process_images(uploaded_files):
-    all_detections = []
-    result_images = []
-    summary_data = []
-    confidences = []
-    temp_dir = tempfile.gettempdir()
-
-    for uploaded_file in uploaded_files:
-        if isinstance(uploaded_file, str):  # Check if it's a sample image path
-            image_path = uploaded_file
-            image = Image.open(image_path)
-        else:
-            image = Image.open(uploaded_file)
-            image_path = os.path.join(temp_dir, f"{uploaded_file.name}")
-            image.save(image_path)
-
-        st.write(f"Detecting in {os.path.basename(image_path)}...")
-        with st.spinner('Running detection...'):
-            result_image, detection_metadata = run(image_path)
-
-        if result_image is not None:
-            result_images.append((result_image, os.path.basename(image_path)))
-            all_detections.append(detection_metadata)
-
-            summary_data.append({
-                "image_name": os.path.basename(image_path),
-                "fish_detected": detection_metadata["fish_count"] > 0,
-                "fish_count": detection_metadata["fish_count"]
-            })
-
-            confidences.extend(detection_metadata["confidences"])
-
-            # Display fish status
-            fish_detected = detection_metadata['fish_count'] > 0
-            fish_status = f"<b><span style='color: green; font-size: 24px;'>YES</span></b> 🐟" if fish_detected else f"<b><span style='color: red; font-size: 24px;'>NO</span></b>"
-
-            st.markdown(f"**Summary for {os.path.basename(image_path)}:** Fish detected: {fish_status}", unsafe_allow_html=True)
-
-            # Display images side by side
-            col1, col2 = st.columns(2)
-            with col1:
-                st.image(image, caption=f"Uploaded Image - {os.path.basename(image_path)}", use_column_width=True)
-            with col2:
-                st.image(result_image, caption=f"Detection Results - {os.path.basename(image_path)}", use_column_width=True)
-
-            st.success(f"Detection completed for {os.path.basename(image_path)} successfully! 🐟")
-
-        else:
-            st.warning(f"No marine ecosystems detected in {os.path.basename(image_path)}.")
-
-    st.session_state["all_detections"] = all_detections
-    return summary_data, confidences
-
-
-# Function to display a summary table and scatter plot side by side with image labels
-def display_summary(summary_data, confidences):
-    if summary_data:
-        df = pd.DataFrame(summary_data)
-
-        col1, col2 = st.columns(2)
-
-        with col1:
-            st.subheader("Summary of Detections")
-            st.table(df[["image_name", "fish_count"]])
-
-        with col2:
-            st.subheader("Fish Detection Confidence Levels")
-            fig, ax = plt.subplots()
-            confidence_index = 0
-
-            for i, row in df.iterrows():
-                num_confidences_for_image = len([c for c in confidences[confidence_index:confidence_index + row["fish_count"]]])
-
-                for j in range(num_confidences_for_image):
-                    if confidence_index < len(confidences):
-                        ax.scatter(confidence_index, confidences[confidence_index], c='blue')
-                        ax.text(confidence_index, confidences[confidence_index], row['image_name'], 
-                                fontsize=10, ha='center', va='bottom', rotation=0)
-                        confidence_index += 1
-
-            ax.axhline(final_confidence, color='red', linestyle='--', label=f'Final Threshold ({final_confidence})')
-            ax.set_xlabel('Detections')
-            ax.set_ylabel('Confidence Level')
-            ax.legend(loc='lower left')
-            st.pyplot(fig)
-
-        if st.session_state.get("all_detections"):
-            json_data = json.dumps(st.session_state["all_detections"], indent=4)
-            st.download_button(
-                label="Download Results as JSON & Reset",
-                data=json_data,
-                file_name="all_detections.json",
-                mime="application/json",
-                key="download_json_bottom"
-            )
-
-# Image uploader with multiple file support
-st.markdown('<div class="custom-file-uploader">', unsafe_allow_html=True)
-uploaded_files = st.file_uploader("Choose image(s)...", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
-
-# Check if files are uploaded, hide the "Auto Run with Sample Images" button if they are
-if not uploaded_files and not st.session_state.get('use_sample_images', False):
-    use_sample_images = st.button("Or Auto Run Using Sample Images", key="sample_button")
-else:
-    use_sample_images = None
-st.markdown('</div>', unsafe_allow_html=True)
-
-# Add the functionality for the "Try it with Sample Images" button
-if use_sample_images:
-    sample_images = load_sample_images()
-    st.session_state['use_sample_images'] = True
-    for sample_image in sample_images:
-        st.session_state.setdefault('uploaded_files', []).append(sample_image)
-    st.session_state['run_automatically'] = True
-
-# Display the Run, Clear, and Download buttons with enhanced styling
-if uploaded_files or st.session_state.get('uploaded_files'):
-    col1, col2, col3 = st.columns([1, 1, 1], gap="small")
-
-    if not st.session_state.get('use_sample_images', False):
-        with col1:
-            run_button = st.button("Click to Run", key="run_button")
-    else:
-        run_button = None
-
-    # Initialize clear_button to None to avoid NameError
-    clear_button = None
-
-    # Conditionally hide the "Clear Results" button while processing
-    with col2:
-        if not st.session_state.get('processing', False):
-            clear_button = st.button("Clear Results", key="clear_button")
-
-    # Run automatically if triggered by the sample images button or the run button
-    if run_button or st.session_state.get('run_automatically'):
-        st.session_state['processing'] = True  # Set the processing flag
-        summary_data, confidences = process_images(uploaded_files or st.session_state['uploaded_files'])
-        display_summary(summary_data, confidences)
-        st.session_state['processing'] = False  # Reset the processing flag after processing is done
-        st.session_state['run_automatically'] = False
-        st.session_state['use_sample_images'] = False
-
-    # Now this check will work, even if clear_button is not defined earlier
-    if clear_button:
-        st.session_state.clear()
-
-    if st.session_state.get("all_detections"):
-        with col3:
-            json_data = json.dumps(st.session_state["all_detections"], indent=4)
-            st.download_button(
-                label="Download Results as JSON & Reset",
-                data=json_data,
-                file_name="all_detections.json",
-                mime="application/json",
-                key="download_json"
-            )
+# Import required modules
+import streamlit as st
+from ultralytics import YOLO
+from PIL import Image
+import os
+import json
+import logging
+import tempfile
+import pandas as pd
+import matplotlib.pyplot as plt
+
+st.set_page_config(
+    page_title="Fish Detector",
+    page_icon="🐟",
+    layout="wide"
+)
+sample_images_folder = "./images/sample_images"
+logging.basicConfig(level=logging.INFO)
+
+# Model loading
+model_folder = "./models"
+st.sidebar.title("🐟 Fish or No Fish Detector")
+st.sidebar.markdown("""
+### For more information:
+- Contact: Michael.Akridge@NOAA.gov
+- Visit the [GitHub repository](https://github.com/MichaelAkridge-NOAA/Fish-or-No-Fish-Detector/)
+""")
+# Display model links
+st.sidebar.markdown("### Model Links")
+st.sidebar.markdown("- [YOLO11 Fish Detector - Grayscale](https://huggingface.co/akridge/yolo11-fish-detector-grayscale)")
+st.sidebar.markdown("- [YOLO11 Segment Fish - Grayscale](https://huggingface.co/akridge/yolo11-segment-fish-grayscale)")
+model_name = st.sidebar.selectbox("Select a YOLO model", os.listdir(model_folder))
+model_path = os.path.join(model_folder, model_name)
+if not os.path.exists(model_path):
+    st.error(f"Model file not found at {model_path}. Please check your setup.")
+    st.stop()
+model = YOLO(model_path)
+
+# Sidebar configuration
+st.sidebar.header("Model Parameters")
+confidence = st.sidebar.slider("Detection Confidence Threshold", 0.0, 1.0, 0.35)
+final_confidence = st.sidebar.slider("Final Yes/No Confidence Threshold", 0.0, 1.0, 0.5)
+
+# Title and description
+st.title("🐟 Fish or No Fish Detector (grayscale)")
+st.write("""
+Is there a fish 🐟 or not? Upload one or more grayscale images to detect fish. Using a trained [Ultralytics YOLO11 Model](https://github.com/ultralytics/ultralytics) for its object detection capabilities.
+
+""")
+
+# Custom CSS for button and uploader alignment
+st.markdown("""
+    <style>
+    .custom-file-uploader {
+        display: flex;
+        align-items: center;
+        margin-top: -10px; /* Adjust to move button closer */
+        justify-content: flex-start;
+    }
+    .css-1cpxqw2 {
+        flex-grow: 1;  /* Let file uploader take remaining space */
+    }
+    .sample-button {
+        font-size: 14px;
+        padding: 8px;
+        background-color: #007BFF;
+        color: white;
+        border: none;
+        border-radius: 5px;
+        cursor: pointer;
+        margin-left: 10px;
+        height: 38px; /* Ensure button matches uploader height */
+    }
+    .sample-button:hover {
+        background-color: #0056b3;
+    }
+    </style>
+""", unsafe_allow_html=True)
+
+# Custom CSS for default button styling
+st.markdown("""
+    <style>
+    .stButton>button, .stDownloadButton>button {
+        width: 100%;
+        padding: 10px;
+        border-radius: 5px;
+        font-size: 18px;
+        font-weight: bold;
+        background-color: #007BFF;
+        color: white;
+        border: none;
+        cursor: pointer;
+    }
+    .stButton>button:hover, .stDownloadButton>button:hover {
+        background-color: #0056b3;
+    }
+    </style>
+""", unsafe_allow_html=True)
+# Load sample images function
+def load_sample_images():
+    return [os.path.join(sample_images_folder, img) for img in os.listdir(sample_images_folder) if img.lower().endswith(('png', 'jpg', 'jpeg'))]
+
+# Prediction function
+def run(image_path):
+    results = model.predict(image_path, conf=confidence)
+    boxes = []
+    fish_count = 0
+    confidences = []
+
+    for box in results[0].boxes:
+        x1, y1, x2, y2 = box.xyxy[0].tolist()
+        conf = box.conf[0].item()
+        class_id = int(box.cls[0].item())
+        class_label = model.names[class_id].lower()  # Normalize to lowercase
+
+        if class_label == "fish" and conf > confidence:
+            fish_count += 1
+        confidences.append(conf)
+
+        boxes.append({"x1": x1, "y1": y1, "x2": x2, "y2": y2, "confidence": conf, "class_id": class_id, "class_label": class_label})
+
+    return results[0].plot()[:, :, ::-1], {"fish_count": fish_count, "confidences": confidences}
+
+# Process images function with directory creation
+# Reusable function to handle multiple image uploads and display results
+def process_images(uploaded_files):
+    all_detections = []
+    result_images = []
+    summary_data = []
+    confidences = []
+    temp_dir = tempfile.gettempdir()
+
+    for uploaded_file in uploaded_files:
+        if isinstance(uploaded_file, str):  # Check if it's a sample image path
+            image_path = uploaded_file
+            image = Image.open(image_path)
+        else:
+            image = Image.open(uploaded_file)
+            image_path = os.path.join(temp_dir, f"{uploaded_file.name}")
+            image.save(image_path)
+
+        st.write(f"Detecting in {os.path.basename(image_path)}...")
+        with st.spinner('Running detection...'):
+            result_image, detection_metadata = run(image_path)
+
+        if result_image is not None:
+            result_images.append((result_image, os.path.basename(image_path)))
+            all_detections.append(detection_metadata)
+
+            summary_data.append({
+                "image_name": os.path.basename(image_path),
+                "fish_detected": detection_metadata["fish_count"] > 0,
+                "fish_count": detection_metadata["fish_count"]
+            })
+
+            confidences.extend(detection_metadata["confidences"])
+
+            # Display fish status
+            fish_detected = detection_metadata['fish_count'] > 0
+            fish_status = f"<b><span style='color: green; font-size: 24px;'>YES</span></b> 🐟" if fish_detected else f"<b><span style='color: red; font-size: 24px;'>NO</span></b>"
+
+            st.markdown(f"**Summary for {os.path.basename(image_path)}:** Fish detected: {fish_status}", unsafe_allow_html=True)
+
+            # Display images side by side
+            col1, col2 = st.columns(2)
+            with col1:
+                st.image(image, caption=f"Uploaded Image - {os.path.basename(image_path)}", use_column_width=True)
+            with col2:
+                st.image(result_image, caption=f"Detection Results - {os.path.basename(image_path)}", use_column_width=True)
+
+            st.success(f"Detection completed for {os.path.basename(image_path)} successfully! 🐟")
+
+        else:
+            st.warning(f"No marine ecosystems detected in {os.path.basename(image_path)}.")
+
+    st.session_state["all_detections"] = all_detections
+    return summary_data, confidences
+
+
+# Function to display a summary table and scatter plot side by side with image labels
+def display_summary(summary_data, confidences):
+    if summary_data:
+        df = pd.DataFrame(summary_data)
+
+        col1, col2 = st.columns(2)
+
+        with col1:
+            st.subheader("Summary of Detections")
+            st.table(df[["image_name", "fish_count"]])
+
+        with col2:
+            st.subheader("Fish Detection Confidence Levels")
+            fig, ax = plt.subplots()
+            confidence_index = 0
+
+            for i, row in df.iterrows():
+                num_confidences_for_image = len([c for c in confidences[confidence_index:confidence_index + row["fish_count"]]])
+
+                for j in range(num_confidences_for_image):
+                    if confidence_index < len(confidences):
+                        ax.scatter(confidence_index, confidences[confidence_index], c='blue')
+                        ax.text(confidence_index, confidences[confidence_index], row['image_name'], 
+                                fontsize=10, ha='center', va='bottom', rotation=0)
+                        confidence_index += 1
+
+            ax.axhline(final_confidence, color='red', linestyle='--', label=f'Final Threshold ({final_confidence})')
+            ax.set_xlabel('Detections')
+            ax.set_ylabel('Confidence Level')
+            ax.legend(loc='lower left')
+            st.pyplot(fig)
+
+        if st.session_state.get("all_detections"):
+            json_data = json.dumps(st.session_state["all_detections"], indent=4)
+            st.download_button(
+                label="Download Results as JSON & Reset",
+                data=json_data,
+                file_name="all_detections.json",
+                mime="application/json",
+                key="download_json_bottom"
+            )
+
+# Image uploader with multiple file support
+st.markdown('<div class="custom-file-uploader">', unsafe_allow_html=True)
+uploaded_files = st.file_uploader("Choose image(s)...", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
+
+# Check if files are uploaded, hide the "Auto Run with Sample Images" button if they are
+if not uploaded_files and not st.session_state.get('use_sample_images', False):
+    use_sample_images = st.button("Or Auto Run Using Sample Images", key="sample_button")
+else:
+    use_sample_images = None
+st.markdown('</div>', unsafe_allow_html=True)
+
+# Add the functionality for the "Try it with Sample Images" button
+if use_sample_images:
+    sample_images = load_sample_images()
+    st.session_state['use_sample_images'] = True
+    for sample_image in sample_images:
+        st.session_state.setdefault('uploaded_files', []).append(sample_image)
+    st.session_state['run_automatically'] = True
+
+# Display the Run, Clear, and Download buttons with enhanced styling
+if uploaded_files or st.session_state.get('uploaded_files'):
+    col1, col2, col3 = st.columns([1, 1, 1], gap="small")
+
+    if not st.session_state.get('use_sample_images', False):
+        with col1:
+            run_button = st.button("Click to Run", key="run_button")
+    else:
+        run_button = None
+
+    # Initialize clear_button to None to avoid NameError
+    clear_button = None
+
+    # Conditionally hide the "Clear Results" button while processing
+    with col2:
+        if not st.session_state.get('processing', False):
+            clear_button = st.button("Clear Results", key="clear_button")
+
+    # Run automatically if triggered by the sample images button or the run button
+    if run_button or st.session_state.get('run_automatically'):
+        st.session_state['processing'] = True  # Set the processing flag
+        summary_data, confidences = process_images(uploaded_files or st.session_state['uploaded_files'])
+        display_summary(summary_data, confidences)
+        st.session_state['processing'] = False  # Reset the processing flag after processing is done
+        st.session_state['run_automatically'] = False
+        st.session_state['use_sample_images'] = False
+
+    # Now this check will work, even if clear_button is not defined earlier
+    if clear_button:
+        st.session_state.clear()
+
+    if st.session_state.get("all_detections"):
+        with col3:
+            json_data = json.dumps(st.session_state["all_detections"], indent=4)
+            st.download_button(
+                label="Download Results as JSON & Reset",
+                data=json_data,
+                file_name="all_detections.json",
+                mime="application/json",
+                key="download_json"
+            )