Update app.py

app.py

@@ -1,41 +1,50 @@
 import streamlit as st
-import kagglehub
-import tensorflow as tf
+from ultralytics import YOLO
 from PIL import Image
-import numpy as np
+import os
 
-# Download model from Kaggle Hub
-model_path = kagglehub.model_download("sohiebalwedyan/seatbelt-detection/tensorFlow2/seatbelt-detection")
-print("Path to model files:", model_path)
-
-# Load the TensorFlow model
-model = tf.saved_model.load(model_path)
+# Load the trained YOLOv8 model
+model = YOLO("yolov8n.pt")
 
+# Define the prediction function
 def predict(image):
-    # Preprocess image
-    img = image.resize((224, 224))
-    img_array = np.array(img) / 255.0
-    img_array = np.expand_dims(img_array, axis=0)
-    
-    # Run inference
-    predictions = model(tf.convert_to_tensor(img_array, dtype=tf.float32))
-    
-    # Process predictions (modify based on actual model output)
-    # This is a placeholder and may need adjustment based on model specifics
-    st.write("Raw predictions:", predictions)
-    
-    return image  # Placeholder return
+    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
-st.title("Seatbelt Detection")
-st.markdown("Upload an image for seatbelt detection.")
+# Streamlit UI for Helmet Detection with YOLO
+st.title("Helmet Detection with YOLOv8")
+st.markdown("Upload an image to detect helmets.")
 
+# 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 = predict(image)
-    # Note: You may need to modify visualization based on model output
+    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)