app.py

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

# Load the YOLOv9c models
detection_model = YOLO('detection.pt')  # Detection model
segmentation_model = YOLO('segmentation.pt')  # Segmentation model

st.title('Protein Crystallization Detection and Segmentation')
st.write('Upload an image to detect or segment crystals.')

# Add a selectbox for mode selection
mode = st.selectbox(
    "Choose the model mode",
    ("Detection", "Segmentation")
)

# Add notes for each mode
if mode == "Detection":
    st.info("Detection is useful for images with distinct and separated crystals.")
else:
    st.info("Segmentation is better for images with many crystals, especially overlaying ones.")

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

if uploaded_file is not None:
    # Read the uploaded image
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption='Uploaded Image', use_column_width=True)

    # Choose the model based on the selected mode
    if mode == "Detection":
        model = detection_model
        st.write("Running Detection...")
    else:
        model = segmentation_model
        st.write("Running Segmentation...")

    # Predict
    results = model.predict(image)

    # Get the results and save the image with bounding boxes
    results_image = results[0].plot()  # This method draws the bounding boxes on the image

    # Convert the results image to a format that Streamlit can display
    results_image = Image.fromarray(results_image)

    # Display the image with detected bounding boxes or masks
    st.image(results_image, caption='Processed Image', use_column_width=True)

    # Display inference time and results in an expander
    with st.expander("Show Results", expanded=True):
        inference_time = (results[0].speed['inference']) / 1000
        st.write(f"Inference Time: {inference_time:.2f} s")
        
        st.write("Detection Results:" if mode == "Detection" else "Segmentation Results:")
        for i, box in enumerate(results[0].boxes):
            x1, y1, x2, y2 = box.xyxy[0]
            conf = box.conf[0]
            cls = box.cls[0]
            st.write(f"{results[0].names[int(cls)]} {i + 1}:")
            st.write(f"  - Bounding Box: ({x1:.2f}, {y1:.2f}, {x2:.2f}, {y2:.2f})")
            st.write(f"  - Confidence: {conf:.2f}")