Update app.py

app.py

@@ -1,50 +1,54 @@
 import streamlit as st
+import torch
 from ultralytics import YOLO
 from PIL import Image
 import os
+import sys
 
-# Load the trained YOLOv8 model for seatbelt detection
-model = YOLO("best.pt")  # Assumes you have a seatbelt-specific trained model
-
-# 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)
-
-# Get example images from the images folder
-def get_example_images():
-    examples = []
-    image_folder = "images"
-    for filename in os.listdir(image_folder):
-        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
-            examples.append(os.path.join(image_folder, filename))
-    return examples
-
-# Streamlit UI for Seatbelt Detection with YOLO
-st.title("Seatbelt Detection with YOLOv8")
-st.markdown("Upload an image to detect seatbelt usage.")
-
-# 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)
-
-# Optionally, show example images from the folder
-if st.checkbox('Show example images'):
-    example_images = get_example_images()
-    for example_image in example_images:
-        img = Image.open(example_image)
-        st.image(img, caption=os.path.basename(example_image), use_column_width=True)
+# 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()}")