Update app.py

app.py

@@ -1,15 +1,16 @@
 import streamlit as st
-from ultralytics import YOLO
 from PIL import Image
+import torch
 import os
 
-# Load the trained YOLOv8 model
-model = YOLO("best.pt")
+# Load the trained YOLOv5 model
+model = torch.hub.load("ultralytics/yolov5", "custom", path="best.pt")  # Load custom 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
+    # Perform inference on the uploaded image
+    results = model(image)  # Runs YOLOv5 model on the uploaded image
+    results_img = results.render()[0]  # Get image with bounding boxes drawn
     return Image.fromarray(results_img)
 
 # Get example images from the images folder
@@ -21,8 +22,8 @@ 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")
+# Streamlit UI for Helmet Detection with YOLOv5
+st.title("Helmet Detection with YOLOv5")
 st.markdown("Upload an image to detect helmets.")
 
 # Allow the user to upload an image