Dogbot/redball_detection.py

63 lines
2.1 KiB
Python

import cv2
import numpy as np
def is_circular(contour, radius):
"""Check if the contour is circular enough to be considered a ball."""
area = cv2.contourArea(contour)
if area <= 0:
return False
circularity = 4 * np.pi * area / (2 * np.pi * radius) ** 2
# Circular enough if circularity is close to 1; adjust threshold as needed
return 0.7 <= circularity <= 1.3
def find_red_ball(image):
# Convert the image to HSV color space
hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
# Adjusted Red color range in HSV; might need fine-tuning
lower_red1 = np.array([0, 120, 70])
upper_red1 = np.array([10, 255, 255])
lower_red2 = np.array([170, 120, 70])
upper_red2 = np.array([180, 255, 255])
mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
# Combine the masks
mask = mask1 + mask2
# Find contours in the mask
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
for contour in contours:
(x, y), radius = cv2.minEnclosingCircle(contour)
if radius > 5 and is_circular(contour, radius): # Checks size and circularity
center = (int(x), int(y))
# Draw the circle around the detected ball
cv2.circle(image, center, int(radius), (0, 255, 0), 2)
# Label the detected ball
cv2.putText(image, 'Red Ball', (center[0] - 20, center[1] - int(radius) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
break # Assuming only one ball, break after finding the first circular object
return image
# Start capturing video from the webcam
cap = cv2.VideoCapture(0)
while True:
_, frame = cap.read()
# Apply the red ball detection and marking
detected_frame = find_red_ball(frame)
# Display the frame with the detected ball
cv2.imshow('Red Ball Detection', detected_frame)
# Break the loop if any key is pressed
if cv2.waitKey(1) != -1:
break
# Release the video capture object and close all OpenCV windows
cap.release()
cv2.destroyAllWindows()