Update app.py

app.py

@@ -3,8 +3,8 @@ from ultralytics import YOLO
 from PIL import Image
 import os
 
-# Load the trained YOLOv8 model
-model = YOLO("best.pt")
+# 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):
@@ -21,9 +21,9 @@ def get_example_images():
             examples.append(os.path.join(image_folder, filename))
     return examples
 
-# Streamlit UI for Helmet Detection with YOLO
-st.title("Helmet Detection with YOLOv8")
-st.markdown("Upload an image to detect helmets.")
+# 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"])
@@ -31,10 +31,10 @@ uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "pn
 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)
@@ -47,4 +47,4 @@ 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)
+        st.image(img, caption=os.path.basename(example_image), use_column_width=True)