Dogbot/redball_chase.py

133 lines
4.7 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
from sensor_msgs.msg import LaserScan
# Define the target stopping distance from the ball in meters
TARGET_DISTANCE = 1.0
# Global variables for state management
current_state = "searching"
ball_position = 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 lidar_callback(data):
"""Process LiDAR data to adjust DogBot's behavior based on distance to obstacles."""
global current_state
if current_state == "approaching":
# Assuming data.ranges is a list of distance measurements
# You might need to process this data differently based on your LiDAR setup
distance_to_ball = min(data.ranges)
if distance_to_ball <= TARGET_DISTANCE:
current_state = "maneuvering"
# Placeholder for function to initiate maneuvering behind the ball
execute_half_circle()
def image_callback(msg):
"""Process image data to detect the red ball and control DogBot's movement."""
global current_state, ball_position
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
twist_msg = Twist()
if center:
x, _ = center
err = x - cv_image.shape[1] / 2
if current_state == "searching" or current_state == "approaching":
current_state = "approaching"
if abs(err) > 100: # Need to rotate DogBot to align with the ball
twist_msg.angular.z = -0.005 * err
else: # Move forward towards the ball
twist_msg.linear.x = -0.2
else:
if current_state == "searching":
twist_msg.angular.z = 0.5 # Rotate DogBot if the ball is not found
cmd_vel_pub.publish(twist_msg)
def execute_half_circle():
"""Executes a half-circle maneuver to position the robot behind the ball."""
global current_state
pub = rospy.Publisher('/cmd_vel', Twist, queue_size=1)
twist_msg = Twist()
# Set angular velocity for a half-circle rotation
twist_msg.angular.z = 0.5 # Positive for counter-clockwise, negative for clockwise
# Set a small or zero linear velocity to pivot around the ball
twist_msg.linear.x = -0.05
# Calculate duration needed to perform a half-circle maneuver around the ball
duration = 2 # Duration in seconds, adjust based on your robot's performance
# Start the maneuver
start_time = rospy.Time.now().to_sec()
while rospy.Time.now().to_sec() - start_time < duration:
pub.publish(twist_msg)
rospy.sleep(0.1) # Sleep briefly to allow message processing
# Stop the robot after completing the maneuver
twist_msg.angular.z = 0
twist_msg.linear.x = 0
pub.publish(twist_msg)
# Update the state to indicate the robot is now positioned behind the ball and ready to push it
current_state = "pushing"
def shutdown():
"""Ensure DogBot stops moving when the script is terminated."""
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.Subscriber("/scan", LaserScan, lidar_callback) # Subscribe to LiDAR data
rospy.on_shutdown(shutdown)
cv2.namedWindow("Red Ball Detection", cv2.WINDOW_NORMAL)
try:
rospy.spin()
except KeyboardInterrupt:
rospy.loginfo("Red ball follower node terminated.")