Update app.py

app.py

@@ -1,16 +1,18 @@
-import cv2
 import torch
+import cv2
 import streamlit as st
 from matplotlib import pyplot as plt
 import numpy as np
+from PIL import Image
+from io import BytesIO
 
 # Function to detect circular objects in the image
-def detect_circles(image_path):
+def detect_circles(image):
     # Load pre-trained YOLOv5 model
     model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # using a smaller model for faster inference
 
-    # Read the image
-    img = cv2.imread(image_path)
+    # Convert PIL image to NumPy array
+    img = np.array(image)
 
     # Run the YOLOv5 model on the image
     results = model(img)
@@ -32,10 +34,11 @@ image_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
 
 if image_file is not None:
     # Open the uploaded image using PIL (Pillow)
-    st.image(image_file, caption="Uploaded Image", use_column_width=True)
+    image = Image.open(image_file)
+    st.image(image, caption="Uploaded Image", use_column_width=True)
 
     # Process the image
-    circles_count, results = detect_circles(image_file)
+    circles_count, results = detect_circles(image)
 
     # Display result
     st.write(f"Number of circular objects detected: {circles_count}")