yolo_model.py

# YOLO model
from ultralytics import YOLO
# Streamlit
import streamlit as st

@st.cache_resource
def load_yolo_model():
    return YOLO("models/best.pt")

def get_detected_objects(yolo_model, image_path, conf_threshold=0.5):
    """
    Run YOLO prediction on an image and return detected objects as a string.
    
    Parameters:
        model_path (str): Path to the YOLO model file.
        image_path (str): Path to the input image.
        conf_threshold (float): Confidence threshold for detections.
        
    Returns:
        str: A comma-separated string of detected object names.
    """
    # Load the YOLO model
    model = yolo_model

    # Run prediction
    results = model.predict(source=image_path, conf=conf_threshold)

    # Extract detected objects as a list
    detected_objects = [box.cls for box in results[0].boxes]  # Access the first image's detections

    # Convert class indices to class names
    detected_class_names = [model.names[int(cls)] for cls in detected_objects]

    # Join detected class names into a single string
    return ", ".join(detected_class_names)