XTDrone/control/yolo_human_tracking.py

46 lines
1.4 KiB
Python
Raw Normal View History

2020-03-23 16:09:44 +08:00
import rospy
from geometry_msgs.msg import Twist
2020-05-20 13:30:28 +08:00
from std_msgs.msg import String
2020-03-23 16:09:44 +08:00
import sys
sys.path.append('/home/robin/catkin_ws/devel/lib/python2.7/dist-packages')
from darknet_ros_msgs.msg import BoundingBoxes
import time
import math
2020-05-20 13:30:28 +08:00
2020-03-23 16:09:44 +08:00
def darknet_callback(data):
2020-05-20 13:30:28 +08:00
global twist, cmd
2020-07-17 11:11:51 +08:00
find = False
2020-03-23 16:09:44 +08:00
for target in data.bounding_boxes:
if(target.id==0):
2020-05-20 13:30:28 +08:00
print('find human')
2020-03-23 16:09:44 +08:00
x_error=y_center-(target.ymax+target.ymin)/2
y_error=x_center-(target.xmax+target.xmin)/2
twist.linear.x = Kp_linear*x_error
twist.angular.z = Kp_angular*math.atan(y_error/x_error)
2020-05-20 13:30:28 +08:00
cmd = ''
2020-07-17 11:11:51 +08:00
find = True
if not find:
twist.linear.x = 0.0
twist.angular.z = 0.0
cmd = 'HOVER'
2020-05-20 13:30:28 +08:00
if __name__ == "__main__":
twist = Twist()
cmd = String()
2020-07-17 11:11:51 +08:00
Kp_linear=0.006
2020-05-20 13:30:28 +08:00
Kp_angular=0.2/math.pi
x_center=752/2
y_center=480/2
vehicle_type = sys.argv[1]
rospy.init_node('yolo_human_tracking')
rospy.Subscriber("/darknet_ros/bounding_boxes", BoundingBoxes, darknet_callback)
vel_pub = rospy.Publisher('/xtdrone/'+vehicle_type+'_0/cmd_vel_flu', Twist, queue_size=2)
cmd_pub = rospy.Publisher('/xtdrone/'+vehicle_type+'_0/cmd', String, queue_size=2)
rate = rospy.Rate(60)
while(True):
rate.sleep()
vel_pub.publish(twist)
cmd_pub.publish(cmd)