Update app.py

app.py

@@ -1,43 +1,43 @@
 import streamlit as st
-import cv2
-import numpy as np
 from PIL import Image
+import torch
+import numpy as np
+import cv2
 
-def detect_proper_circles(image):
-    # Convert to grayscale
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-
-    # Apply adaptive thresholding to detect edges clearly
-    thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
-                                   cv2.THRESH_BINARY_INV, 11, 2)
-
-    # Find contours in the thresholded image
-    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    proper_circles = 0
-
-    for contour in contours:
-        # Fit an enclosing circle
-        center, radius = cv2.minEnclosingCircle(contour)
-        center = (int(center[0]), int(center[1]))
-        radius = int(radius)
-
-        # Calculate circularity: Area of contour vs. area of enclosing circle
-        area = cv2.contourArea(contour)
-        enclosing_area = np.pi * (radius ** 2)
-        circularity = area / enclosing_area if enclosing_area > 0 else 0
-
-        # Set circularity and radius thresholds for proper circles
-        if circularity > 0.75 and radius > 10:  # 0.75 ensures near-perfect circles
-            # Draw the detected circle
-            cv2.circle(image, center, radius, (0, 255, 0), 2)
-            proper_circles += 1
-
-    return image, proper_circles
-
-# Streamlit App
+# Load the trained YOLOv5 model
+@st.cache_resource
+def load_model():
+    model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')
+    return model
+
+# Circle detection function
+def detect_circles(model, image):
+    # Convert PIL Image to OpenCV format
+    image_cv = np.array(image)
+    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)
+    
+    # Run inference
+    results = model(image_cv)
+    
+    # Extract detections and count
+    detections = results.pandas().xyxy[0]
+    circle_count = len(detections)
+    
+    # Draw bounding boxes
+    for _, row in detections.iterrows():
+        x1, y1, x2, y2, conf = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax']), row['confidence']
+        cv2.rectangle(image_cv, (x1, y1), (x2, y2), (0, 255, 0), 2)
+        cv2.putText(image_cv, f"Circle {conf:.2f}", (x1, y1 - 10), 
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+
+    return image_cv, circle_count
+
+# Streamlit app
 def main():
-    st.title("Proper Circle Detector and Counter")
-    st.write("Upload an image to detect and count proper circular shapes like pipes or log ends.")
+    st.title("Circle Detector and Counter")
+    st.write("Upload an image to detect and count circular objects like logs or pipes.")
+
+    model = load_model()  # Load YOLOv5 model
 
     uploaded_file = st.file_uploader("Upload Image", type=["jpg", "png", "jpeg"])
     if uploaded_file is not None:
@@ -45,17 +45,13 @@ def main():
         image = Image.open(uploaded_file)
         st.image(image, caption="Uploaded Image", use_column_width=True)
 
-        # Convert to OpenCV format
-        image_cv = np.array(image)
-        image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)
-
-        # Detect proper circles
-        processed_image, circle_count = detect_proper_circles(image_cv)
+        # Detect circles
+        processed_image, circle_count = detect_circles(model, image)
 
         # Display results
         st.image(cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB), 
                  caption=f"Processed Image - Circles Counted: {circle_count}", use_column_width=True)
-        st.success(f"Number of proper circles detected: {circle_count}")
+        st.success(f"Number of circles detected: {circle_count}")
 
 if __name__ == "__main__":
     main()