app.py

import streamlit as st
import torch
from ultralytics import YOLO
from PIL import Image
import os
import sys

# Print Python and library versions for debugging
st.write(f"Python version: {sys.version}")
st.write(f"Torch version: {torch.__version__}")
st.write(f"Current working directory: {os.getcwd()}")
st.write(f"Files in current directory: {os.listdir('.')}")

# Check if the model file exists
model_path = "best.pt"  # or the path to your actual model file
if not os.path.exists(model_path):
    st.error(f"Model file {model_path} not found!")
else:
    st.success(f"Model file {model_path} found!")

try:
    # Load the trained YOLOv8 model
    model = YOLO(model_path)

    # Define the prediction function
    def predict(image):
        results = model(image)  # Run YOLOv8 model on the uploaded image
        results_img = results[0].plot()  # Get image with bounding boxes
        return Image.fromarray(results_img)

    # Streamlit UI for Object Detection
    st.title("Object Detection with YOLOv8")
    st.markdown("Upload an image for detection.")

    # Allow the user to upload an image
    uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    if uploaded_image is not None:
        # Open the uploaded image using PIL
        image = Image.open(uploaded_image)
        
        # Display the uploaded image
        st.image(image, caption="Uploaded Image", use_column_width=True)
        
        # Run the model prediction
        st.subheader("Prediction Results:")
        result_image = predict(image)
        
        # Display the result image with bounding boxes
        st.image(result_image, caption="Detected Image", use_column_width=True)

except Exception as e:
    st.error(f"An error occurred: {e}")
    st.error(f"Traceback: {sys.exc_info()}")