Update app.py

app.py

@@ -1,54 +1,37 @@
 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()}")
+# Directly load the model using full path if needed
+model_path = os.path.join(os.getcwd(), 'best.pt')
+
+# Load the YOLO 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)