Update app.py

app.py

@@ -1,68 +1,48 @@
+import cv2
+import torch
 import streamlit as st
+from matplotlib import pyplot as plt
 import numpy as np
-import cv2
-from PIL import Image
-
-# Function to process the image and count coins
-def count_coins(image):
-    # Convert PIL Image to OpenCV format
-    img = np.array(image)
-    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-    
-    # Resize image
-    img = cv2.resize(img, (640, 800))
-    image_copy = img.copy()
 
-    # Apply Gaussian Blur
-    img = cv2.GaussianBlur(img, (7, 7), 3)
+# Function to detect circular objects in the image
+def detect_circles(image_path):
+    # Load pre-trained YOLOv5 model
+    model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # using a smaller model for faster inference
 
-    # Convert to grayscale
-    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    # Read the image
+    img = cv2.imread(image_path)
 
-    # Apply thresholding
-    ret, thresh = cv2.threshold(gray, 170, 255, cv2.THRESH_BINARY)
-
-    # Find contours
-    contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
-
-    # Calculate contour areas
-    area = {}
-    for i in range(len(contours)):
-        cnt = contours[i]
-        ar = cv2.contourArea(cnt)
-        area[i] = ar
-
-    # Sort contours by area
-    srt = sorted(area.items(), key=lambda x: x[1], reverse=True)
-    results = np.array(srt).astype("int")
+    # Run the YOLOv5 model on the image
+    results = model(img)
     
-    # Count coins (objects with area greater than a threshold)
-    num = np.argwhere(results[:, 1] > 500).shape[0]
+    # Filter out circular objects (based on class label, you may adjust this depending on what YOLO detects)
+    circles = 0
+    for det in results.xywh[0]:
+        # Class 44 is for "circular" items like wheels, stop signs, etc. (may vary by model and training)
+        if int(det[5]) == 44:  # Replace 44 with the appropriate class ID for circular objects
+            circles += 1
+
+    return circles, results
 
-    # Draw contours around coins
-    for i in range(1, num):
-        image_copy = cv2.drawContours(image_copy, contours, results[i, 0], (0, 255, 0), 3)
+# Streamlit UI
+st.title("Circular Object Detection with YOLOv5")
 
-    return image_copy, num - 1  # Return processed image and the count of coins
+# Upload image file
+image_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
 
-# Streamlit app
-def main():
-    st.title("Coin Detection and Counting")
-    st.write("Upload an image to detect and count coins.")
+if image_file is not None:
+    # Open the uploaded image using PIL (Pillow)
+    st.image(image_file, caption="Uploaded Image", use_column_width=True)
 
-    # File uploader for image
-    uploaded_file = st.file_uploader("Upload Image", type=["jpg", "png", "jpeg"])
-    if uploaded_file is not None:
-        # Open the uploaded image with PIL
-        image = Image.open(uploaded_file)
-        st.image(image, caption="Uploaded Image", use_column_width=True)
+    # Process the image
+    circles_count, results = detect_circles(image_file)
 
-        # Count coins
-        processed_image, coin_count = count_coins(image)
+    # Display result
+    st.write(f"Number of circular objects detected: {circles_count}")
 
-        # Display processed image and the result
-        st.image(cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB), caption=f"Processed Image - Coins Counted: {coin_count}", use_column_width=True)
-        st.success(f"Number of coins detected: {coin_count}")
+    # Display detection output
+    fig, ax = plt.subplots(1, 1, figsize=(10, 6))
+    ax.imshow(results.render()[0])
+    plt.show()
 
-if __name__ == "__main__":
-    main()
+    st.pyplot(fig)