Create app.py

app.py

@@ -0,0 +1,50 @@
+import streamlit as st
+from ultralytics import YOLO
+from PIL import Image
+import os
+
+# Load the trained YOLOv8 model
+model = YOLO("best.pt")
+
+# 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 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_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)