roofyv5 / app.py
scfive's picture
Update app.py
8873f80 verified
raw
history blame
1.15 kB
import streamlit as st
from ultralytics import YOLO
from PIL import Image
import numpy as np
import os
# Define the model path
model_path = "best.pt" # Assumes `best.pt` is in the root directory
# Verify file existence
if not os.path.exists(model_path):
st.error(f"Model file not found at: {model_path}. Current directory files: {os.listdir()}")
# Load the YOLO model
try:
model = YOLO(model_path)
except Exception as e:
st.error(f"Error loading model: {e}")
# Streamlit app
st.title("YOLOv11 Object Detection")
st.write("Upload an image and let the model detect objects.")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
if uploaded_file:
# Read and display the image
image = Image.open(uploaded_file)
st.image(image, caption="Uploaded Image", use_column_width=True)
# Perform prediction
with st.spinner("Processing..."):
try:
results = model.predict(np.array(image))
st.write("Detection Results:")
st.image(results[0].plot(), caption="Detections", use_column_width=True)
except Exception as e:
st.error(f"E