124 lines
4.1 KiB
Python
124 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import rospy
|
|
import cv2
|
|
import numpy as np
|
|
from sensor_msgs.msg import Image
|
|
from geometry_msgs.msg import Twist
|
|
from cv_bridge import CvBridge, CvBridgeError
|
|
import random
|
|
import math
|
|
|
|
# Define the target stopping distance from the ball in pixels
|
|
RADIUS_THRESHOLD = 40
|
|
|
|
# Global variables for state management
|
|
current_state = "searching"
|
|
ball_position = None
|
|
ball_radius = None
|
|
|
|
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
|
|
return 0.7 <= circularity <= 1.3
|
|
|
|
def find_red_ball(image):
|
|
"""Detect red ball in the image and return its position and radius if found."""
|
|
hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
|
|
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)
|
|
mask = mask1 + mask2
|
|
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):
|
|
return (int(x), int(y)), radius
|
|
return None, None
|
|
|
|
def execute_search():
|
|
"""Execute a search pattern: rotate 360, move randomly if the ball isn't found, and repeat."""
|
|
global current_state
|
|
pub = rospy.Publisher('/cmd_vel', Twist, queue_size=1)
|
|
twist_msg = Twist()
|
|
|
|
# Rotate 360 degrees
|
|
twist_msg.angular.z = 0.5 # Adjust as needed
|
|
rotate_duration = 2 * math.pi / twist_msg.angular.z # Full circle
|
|
start_time = rospy.Time.now().to_sec()
|
|
while rospy.Time.now().to_sec() - start_time < rotate_duration:
|
|
pub.publish(twist_msg)
|
|
rospy.sleep(0.1)
|
|
twist_msg.angular.z = 0
|
|
pub.publish(twist_msg) # Stop rotation
|
|
|
|
# Move in a random direction if the ball isn't found
|
|
if current_state == "searching":
|
|
twist_msg.linear.x = 0.2
|
|
move_duration = random.uniform(1, 3) # Random duration between 1 and 3 seconds
|
|
start_time = rospy.Time.now().to_sec()
|
|
while rospy.Time.now().to_sec() - start_time < move_duration:
|
|
pub.publish(twist_msg)
|
|
twist_msg.linear.x = 0
|
|
pub.publish(twist_msg) # Stop moving
|
|
|
|
def image_callback(msg):
|
|
global current_state, ball_position, ball_radius
|
|
try:
|
|
cv_image = bridge.imgmsg_to_cv2(msg, "bgr8")
|
|
except CvBridgeError as e:
|
|
rospy.logerr(e)
|
|
|
|
center, radius = find_red_ball(cv_image)
|
|
ball_position = center
|
|
ball_radius = radius
|
|
|
|
if center and radius >= RADIUS_THRESHOLD:
|
|
# Ball found and close enough, execute half-circle
|
|
current_state = "maneuvering"
|
|
execute_half_circle()
|
|
elif current_state == "searching":
|
|
# Ball not found or not close enough, continue searching
|
|
execute_search()
|
|
|
|
def execute_half_circle():
|
|
global current_state
|
|
pub = rospy.Publisher('/cmd_vel', Twist, queue_size=1)
|
|
twist_msg = Twist()
|
|
twist_msg.angular.z = 0.5 # Adjust for half-circle
|
|
twist_msg.linear.x = 0.05 # Slight forward movement
|
|
duration = 4 # Adjust for a proper half-circle based on your robot
|
|
start_time = rospy.Time.now().to_sec()
|
|
while rospy.Time.now().to_sec() - start_time < duration:
|
|
pub.publish(twist_msg)
|
|
twist_msg.angular.z = 0
|
|
twist_msg.linear.x = 0
|
|
pub.publish(twist_msg) # Stop
|
|
current_state = "approaching"
|
|
|
|
def shutdown():
|
|
rospy.loginfo("Shutting down")
|
|
cmd_vel_pub.publish(Twist())
|
|
cv2.destroyAllWindows()
|
|
|
|
if __name__ == '__main__':
|
|
rospy.init_node('redball_chase', anonymous=False)
|
|
bridge = CvBridge()
|
|
cmd_vel_pub = rospy.Publisher('/cmd_vel', Twist, queue_size=1)
|
|
rospy.Subscriber("/camera/image", Image, image_callback)
|
|
|
|
rospy.on_shutdown(shutdown)
|
|
cv2.namedWindow("Red Ball Detection", cv2.WINDOW_NORMAL)
|
|
|
|
try:
|
|
rospy.spin()
|
|
except KeyboardInterrupt:
|
|
rospy.loginfo("Red ball follower node terminated.")
|
|
|