pheerasit/yolo-streamlit-test

1import streamlit as st2from ultralytics import YOLO3from PIL import Image4import numpy as np5import cv26 7# --- App Title and Description ---8# st.set_page_config(page_title="Object Detection App", layout="wide")9st.title("🖼️ Object Detection with YOLOv11")10st.write(11    "Upload an image and the YOLOv8 model will detect objects in it. "12    "You can then upload another image to try again."13)14 15# --- Model Loading ---16# Load a pre-trained YOLOv8 model17# 'yolo11n.pt' is a small and fast model, ideal for general purposes.18try:19    model = YOLO('yolo11n.pt')20except Exception as e:21    st.error(f"Error loading model: {e}")22    st.stop()23 24 25# --- Image Uploader ---26uploaded_file = st.file_uploader(27    "Choose an image...",28    type=["jpg", "jpeg", "png"]29)30 31 32if uploaded_file is not None:33    # --- Display Images Side-by-Side ---34    col1, col2 = st.columns(2)35 36    # Open the uploaded image37    original_image = Image.open(uploaded_file)38 39    with col1:40        st.header("Original Image")41        st.image(original_image, caption="Your uploaded image.", use_container_width=True)42 43    with col2:44        st.header("Detected Image")45        with st.spinner('Detecting objects...'):46            # Perform object detection47            results = model(original_image)48 49            # The '.plot()' method returns a NumPy array with bounding boxes drawn on it50            detected_image_np = results[0].plot()51 52            # The result from '.plot()' is in BGR format (used by OpenCV).53            # We need to convert it to RGB for correct display in Streamlit.54            detected_image_rgb = cv2.cvtColor(detected_image_np, cv2.COLOR_BGR2RGB)55 56            st.image(detected_image_rgb, caption="Image with detected objects.", use_container_width=True)57 58    # --- Optional: Display Detection Details ---59    # You can uncomment the following lines to show the detected class names and confidences60    st.header("Detection Details")61    for result in results:62        for box in result.boxes:63            class_id = int(box.cls[0])64            class_name = model.names[class_id]65            confidence = float(box.conf[0])66            st.write(f"- **{class_name}**: {confidence:.2f} confidence")