Update app.py

app.py

@@ -1,37 +1,35 @@
 import streamlit as st
-from ultralytics import YOLO
+import torch
+import cv2
+import numpy as np
 from PIL import Image
-import os
 
-# Directly load the model using full path if needed
-model_path = os.path.join(os.getcwd(), 'best.pt')
+# Load YOLOv5 model
+model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
 
-# Load the YOLO 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)
+    # Convert PIL Image to numpy array
+    img_array = np.array(image)
+    
+    # Detect objects
+    results = model(img_array)
+    
+    # Plot results
+    img_with_boxes = results.render()[0]
+    
+    # Convert back to PIL Image
+    return Image.fromarray(img_with_boxes)
 
-# Streamlit UI for Object Detection
-st.title("Object Detection with YOLOv8")
+# Streamlit UI
+st.title("Object Detection with YOLOv5")
 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)