From c7177c19d21b740235fded9c0cffd2322df3cf26 Mon Sep 17 00:00:00 2001 From: Sameer Parekh Date: Tue, 23 Oct 2012 05:51:14 -0700 Subject: [PATCH] initial commit, moving files over from ardrone_autonomy fork --- .gitignore | 28 +++ CMakeLists.txt | 30 +++ Makefile | 1 + launch/track.launch | 17 ++ mainpage.dox | 14 ++ manifest.xml | 21 ++ nodes/ardrone_control.py | 73 ++++++ nodes/ardrone_follow.py | 261 +++++++++++++++++++++ nodes/ardrone_teleop_joy.py | 62 +++++ nodes/ardrone_tracker.py | 74 ++++++ nodes/pid.py | 88 +++++++ nodes/tracker.py | 449 ++++++++++++++++++++++++++++++++++++ 12 files changed, 1118 insertions(+) create mode 100644 .gitignore create mode 100644 CMakeLists.txt create mode 100644 Makefile create mode 100644 launch/track.launch create mode 100644 mainpage.dox create mode 100644 manifest.xml create mode 100755 nodes/ardrone_control.py create mode 100755 nodes/ardrone_follow.py create mode 100755 nodes/ardrone_teleop_joy.py create mode 100755 nodes/ardrone_tracker.py create mode 100644 nodes/pid.py create mode 100644 nodes/tracker.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..73412b0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,28 @@ +# Compiled Object files +*.slo +*.lo +*.o + +# Compiled Dynamic libraries +*.so + +# Compiled Static libraries +*.lai +*.la +*.a +/build/CMakeFiles/ + +# Compiled python files +*.pyc + +# emacs files +*~ + +# bag files +*.bag + +# specific directories +bin/ +build/ +msg_gen/ +srv_gen/ diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..f8f1c9c --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,30 @@ +cmake_minimum_required(VERSION 2.4.6) +include($ENV{ROS_ROOT}/core/rosbuild/rosbuild.cmake) + +# Set the build type. Options are: +# Coverage : w/ debug symbols, w/o optimization, w/ code-coverage +# Debug : w/ debug symbols, w/o optimization +# Release : w/o debug symbols, w/ optimization +# RelWithDebInfo : w/ debug symbols, w/ optimization +# MinSizeRel : w/o debug symbols, w/ optimization, stripped binaries +#set(ROS_BUILD_TYPE RelWithDebInfo) + +rosbuild_init() + +#set the default path for built executables to the "bin" directory +set(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/bin) +#set the default path for built libraries to the "lib" directory +set(LIBRARY_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/lib) + +#uncomment if you have defined messages +#rosbuild_genmsg() +#uncomment if you have defined services +#rosbuild_gensrv() + +#common commands for building c++ executables and libraries +#rosbuild_add_library(${PROJECT_NAME} src/example.cpp) +#target_link_libraries(${PROJECT_NAME} another_library) +#rosbuild_add_boost_directories() +#rosbuild_link_boost(${PROJECT_NAME} thread) +#rosbuild_add_executable(example examples/example.cpp) +#target_link_libraries(example ${PROJECT_NAME}) diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..b75b928 --- /dev/null +++ b/Makefile @@ -0,0 +1 @@ +include $(shell rospack find mk)/cmake.mk \ No newline at end of file diff --git a/launch/track.launch b/launch/track.launch new file mode 100644 index 0000000..b1b71b9 --- /dev/null +++ b/launch/track.launch @@ -0,0 +1,17 @@ + + + + + + + + + + + + + + + + + diff --git a/mainpage.dox b/mainpage.dox new file mode 100644 index 0000000..371b16e --- /dev/null +++ b/mainpage.dox @@ -0,0 +1,14 @@ +/** +\mainpage +\htmlinclude manifest.html + +\b falkor_ardrone + + + +--> + + +*/ diff --git a/manifest.xml b/manifest.xml new file mode 100644 index 0000000..69780fe --- /dev/null +++ b/manifest.xml @@ -0,0 +1,21 @@ + + + + falkor_ardrone + + + Sameer Parekh + BSD + + http://ros.org/wiki/falkor_ardrone + + + + + + + + + + + diff --git a/nodes/ardrone_control.py b/nodes/ardrone_control.py new file mode 100755 index 0000000..1c80e00 --- /dev/null +++ b/nodes/ardrone_control.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python +import roslib +roslib.load_manifest('falkor_ardrone') + +import rospy +from geometry_msgs.msg import Twist +from std_msgs.msg import Empty +from ardrone_autonomy.msg import Navdata +import pid +import time + +class ArdroneControl: + def __init__( self ): + self.nav_sub = rospy.Subscriber( "ardrone/navdata", Navdata, self.callback_navdata ) + self.cmd_vel_pub = rospy.Publisher( "cmd_vel", Twist ) + self.goal_vel_sub = rospy.Subscriber( "goal_vel", Twist, self.callback_goal_vel ) + + # gain_p, gain_i, gain_d + self.linearxpid = pid.Pid2( 0.5, 0.0, 0.5 ) + self.linearypid = pid.Pid2( 0.5, 0.0, 0.5 ) + + self.vx = self.vy = self.vz = self.ax = self.ay = self.az = 0.0 + self.last_time = None + self.goal_vel = Twist() + + def callback_goal_vel( self, data ): + self.goal_vel = data + + def callback_navdata( self, data ): + self.vx = data.vx/1e3 + self.vy = data.vy/1e3 + self.vz = data.vz/1e3 + + # these accelerations must take into account orientation, so we get acceleration relative + # to base_stabilized. This does not. +# self.ax = (data.ax*9.82) +# self.ay = (data.ay*9.82) +# self.az = (data.az - 1.0)*9.82 + + def update( self ): + if self.last_time == None: + self.last_time = rospy.Time.now() + dt = 0.0 + else: + time = rospy.Time.now() + dt = ( time - self.last_time ).to_sec() + self.last_time = time + + cmd = Twist() + cmd.angular.y = 0 + cmd.angular.x = 0 + cmd.angular.z = self.goal_vel.angular.z + cmd.linear.z = self.goal_vel.linear.z + cmd.linear.x = self.linearxpid.update( self.goal_vel.linear.x, self.vx, 0.0, dt ) + cmd.linear.y = self.linearypid.update( self.goal_vel.linear.y, self.vy, 0.0, dt ) + + self.cmd_vel_pub.publish( cmd ) + +def main(): + rospy.init_node( 'ardrone_control' ) + + control = ArdroneControl() + r = rospy.Rate(100) + + try: + while not rospy.is_shutdown(): + control.update() + r.sleep() + except KeyboardInterrupt: + print "Shutting down" + +if __name__ == '__main__': + main() diff --git a/nodes/ardrone_follow.py b/nodes/ardrone_follow.py new file mode 100755 index 0000000..d1ec2ae --- /dev/null +++ b/nodes/ardrone_follow.py @@ -0,0 +1,261 @@ +#!/usr/bin/env python +import roslib +roslib.load_manifest('falkor_ardrone') + +import sys +import rospy +import math +import pid +import cv2 +import numpy as np +from cv_bridge import CvBridge, CvBridgeError + +from geometry_msgs.msg import Point +from geometry_msgs.msg import Twist +from std_msgs.msg import Empty +from sensor_msgs.msg import Joy, Image +from ardrone_autonomy.msg import Navdata +from ardrone_autonomy.srv import LedAnim + +class ArdroneFollow: + def __init__( self ): + self.led_service = rospy.ServiceProxy( "ardrone/setledanimation", LedAnim ) +# print "waiting for driver to startup" +# rospy.wait_for_service( self.led_service ) + + self.tracker_sub = rospy.Subscriber( "ardrone_tracker/found_point", + Point, self.found_point_cb ) + self.goal_vel_pub = rospy.Publisher( "goal_vel", Twist ) + self.found_time = None + + self.tracker_img_sub = rospy.Subscriber( "ardrone_tracker/image", + Image, self.image_cb ) + self.tracker_image = None + + self.timer = rospy.Timer( rospy.Duration( 0.10 ), self.timer_cb, False ) + + self.land_pub = rospy.Publisher( "ardrone/land", Empty ) + self.takeoff_pub = rospy.Publisher( "ardrone/takeoff", Empty ) + self.reset_pub = rospy.Publisher( "ardrone/reset", Empty ) + + self.angularZlimit = 3.141592 / 2 + self.linearXlimit = 1.0 + self.linearZlimit = 2.0 + + self.xPid = pid.Pid( 0.020, 0.0, 0.0, self.angularZlimit ) + self.yPid = pid.Pid( 0.020, 0.0, 0.0, self.linearZlimit ) + self.zPid = pid.Pid( 0.070, 0.0, 0.0, self.linearXlimit ) + + self.xPid.setPointMin = 40 + self.xPid.setPointMax = 60 + + self.yPid.setPointMin = 40 + self.yPid.setPointMax = 60 + + self.zPid.setPointMin = 17 + self.zPid.setPointMax = 23 + + self.lastAnim = -1 + + self.found_point = Point( 0, 0, -1 ) + self.old_cmd = self.current_cmd = Twist() + + self.joy_sub = rospy.Subscriber( "joy", Joy, self.callback_joy ) + self.manual_cmd = False + self.auto_cmd = False + + self.bridge = CvBridge() + + cv2.namedWindow( 'AR.Drone Follow', cv2.cv.CV_WINDOW_NORMAL ) + + def image_cb( self, data ): + try: + cv_image = self.bridge.imgmsg_to_cv( data, "passthrough" ) + except CvBridgeError, e: + print e + + self.tracker_image = np.asarray( cv_image ) + + def increase_z_setpt( self ): + self.zPid.setPointMin *= 1.01 + self.zPid.setPointMax *= 1.01 + + def decrease_z_setpt( self ): + self.zPid.setPointMin /= 1.01 + self.zPid.setPointMax /= 1.01 + + def callback_joy( self, data ): + empty_msg = Empty() + + if data.buttons[12] == 1 and self.last_buttons[12] == 0: + self.takeoff() + + if data.buttons[14] == 1 and self.last_buttons[14] == 0: + self.land() + + if data.buttons[13] == 1 and self.last_buttons[13] == 0: + self.reset() + + if data.buttons[15] == 1 and self.last_buttons[15] == 0: + self.auto_cmd = not self.auto_cmd + self.hover() + + if data.buttons[4] == 1: + self.increase_z_setpt() + + if data.buttons[6] == 1: + self.decrease_z_setpt() + + self.last_buttons = data.buttons + + # Do cmd_vel + self.current_cmd = Twist() + + self.current_cmd.angular.x = self.current_cmd.angular.y = 0 + self.current_cmd.angular.z = data.axes[2] * self.angularZlimit + + self.current_cmd.linear.z = data.axes[3] * self.linearZlimit + self.current_cmd.linear.x = data.axes[1] * self.linearXlimit + self.current_cmd.linear.y = data.axes[0] * self.linearXlimit + + if ( self.current_cmd.linear.x == 0 and + self.current_cmd.linear.y == 0 and + self.current_cmd.linear.z == 0 and + self.current_cmd.angular.z == 0 ): + self.manual_cmd = False + else: + self.setLedAnim( 9 ) + self.manual_cmd = True + + self.goal_vel_pub.publish( self.current_cmd ) + + def setLedAnim( self, animType, freq = 10 ): + if self.lastAnim == type: + return + + msg = LedAnim(); + msg.type = animType; + msg.freq = freq; + msg.duration = 3600; + + self.led_service( type = animType, freq = freq, duration = 255 ) + self.lastAnim = type + + def takeoff( self ): + self.takeoff_pub.publish( Empty() ) + self.setLedAnim( 9 ) + + def land( self ): + self.land_pub.publish( Empty() ) + + def reset( self ): + self.reset_pub.publish( Empty() ) + + def navdata_cb( self, data ): + self.vx = data.vx/1e3 + self.vy = data.vy/1e3 + self.vz = data.vz/1e3 + self.ax = (data.ax*9.82) + self.ay = (data.ay*9.82) + self.az = (data.az - 1)*9.82 + + def found_point_cb( self, data ): + self.found_point = data + self.found_time = rospy.Time.now() + + def hover( self ): + hoverCmd = Twist() + self.goal_vel_pub.publish( hoverCmd ) + + def hover_cmd_cb( self, data ): + self.hover() + + def draw_picture( self ): + if self.tracker_image == None: + return + + vis = self.tracker_image.copy() + (ySize, xSize, depth) = vis.shape + + cx_min = int(self.xPid.setPointMin * xSize / 100) + cx_max = int(self.xPid.setPointMax * xSize / 100) + cx = int( ( cx_min + cx_max ) / 2 ) + + cy_min = int(self.yPid.setPointMin * ySize / 100) + cy_max = int(self.yPid.setPointMax * ySize / 100) + cy = int( ( cy_min + cy_max ) / 2 ) + + cv2.rectangle( vis, ( cx_min, cy_min ), ( cx_max, cy_max ), + ( 0, 255, 255 ) ) + + areasqrt = np.sqrt( 640 * 480 ) + side_min_half = int( self.zPid.setPointMin * areasqrt / 100 / 2 ) + side_max_half = int( self.zPid.setPointMax * areasqrt / 100 / 2 ) + + cv2.rectangle( vis, ( cx - side_min_half, cy - side_min_half ), + ( cx + side_min_half, cy + side_min_half ), + ( 255, 255, 0 ) ) + + cv2.rectangle( vis, ( cx - side_max_half, cy - side_max_half ), + ( cx + side_max_half, cy + side_max_half ), + ( 255, 255, 0 ) ) + + if self.current_cmd.linear.x > 0: + line_color = ( 0, 255, 0 ) + else: + line_color = ( 0, 0, 255 ) + + cv2.line( vis, ( 320, 180 ), ( int( xSize/2 - self.current_cmd.angular.z * 320 ), + int( ySize/2 - self.current_cmd.linear.z * 180 ) ), + line_color, + min( max( int( abs( self.current_cmd.linear.x ) * 255 ), 0 ), 255 ) ) + + cv2.imshow( 'AR.Drone Follow', vis ) + cv2.waitKey( 1 ) + + def timer_cb( self, event ): + self.draw_picture() + + # If we haven't received a found point in a second, let found_point be + # (0,0,-1) + if ( self.found_time == None or + ( rospy.Time.now() - self.found_time ).to_sec() > 1 ): + self.found_point = Point( 0, 0, -1.0 ) + self.found_time = rospy.Time.now() + + if event.last_real == None: + dt = 0 + else: + dt = ( event.current_real - event.last_real ).to_sec() + + if self.found_point.z == -1.0: + self.current_cmd = Twist() + self.setLedAnim( 0, 2 ) + else: + self.current_cmd.angular.z = self.xPid.get_output( self.found_point.x, + dt ) + self.current_cmd.linear.z = self.yPid.get_output( self.found_point.y, + dt ) + self.current_cmd.linear.x = self.zPid.get_output( self.found_point.z, + dt ) + + self.setLedAnim( 8, 2 ) + + if self.auto_cmd == False or self.manual_cmd == True: + self.setLedAnim( 9 ) + return + + self.goal_vel_pub.publish( self.current_cmd ) + + +def main(): + rospy.init_node( 'ardrone_follow' ) + af = ArdroneFollow() + + try: + rospy.spin() + except KeyboardInterrupt: + print "Keyboard interrupted" + +if __name__ == '__main__': + main() diff --git a/nodes/ardrone_teleop_joy.py b/nodes/ardrone_teleop_joy.py new file mode 100755 index 0000000..2a92a64 --- /dev/null +++ b/nodes/ardrone_teleop_joy.py @@ -0,0 +1,62 @@ +#!/usr/bin/env python +import roslib +roslib.load_manifest('falkor_ardrone') + +import rospy +from geometry_msgs.msg import Twist +from sensor_msgs.msg import Joy +from std_msgs.msg import Empty +from ardrone_autonomy.msg import Navdata +import pid +import time + +class ArdroneTeleopJoy: + + def __init__( self ): + self.cmd_vel_pub = rospy.Publisher( "cmd_vel_interm", Twist ) + self.land_pub = rospy.Publisher( "ardrone/land", Empty ) + self.takeoff_pub = rospy.Publisher( "ardrone/takeoff", Empty ) + self.reset_pub = rospy.Publisher( "ardrone/reset", Empty ) + self.joy_sub = rospy.Subscriber( "joy", Joy, self.callback_joy ) + + def callback_joy( self, data ): + empty_msg = Empty() + + if data.buttons[12] == 1 and self.last_buttons[12] == 0: + self.takeoff_pub.publish( empty_msg ) + self.pause = 5.0 + + if data.buttons[14] == 1 and self.last_buttons[14] == 0: + self.land_pub.publish( empty_msg ) + + if data.buttons[13] == 1 and self.last_buttons[13] == 0: + self.reset_pub.publish( empty_msg ) + + self.last_buttons = data.buttons + + # Do cmd_vel + cmd = Twist() + + cmd.angular.x = cmd.angular.y = 0 + cmd.angular.z = data.axes[2] * 3.14/2 + + cmd.linear.z = data.axes[3] * 2.0 + cmd.linear.x = data.axes[1] * 1.0 + cmd.linear.y = data.axes[0] * 1.0 + self.cmd_vel_pub.publish( cmd ) + + +def main(): + rospy.init_node( 'ardrone_teleop_joy' ) + + teleop = ArdroneTeleopJoy() + + try: + rospy.spin() + except KeyboardInterrupt: + print "Shutting down" + +if __name__ == '__main__': + main() + + diff --git a/nodes/ardrone_tracker.py b/nodes/ardrone_tracker.py new file mode 100755 index 0000000..4f9f8c9 --- /dev/null +++ b/nodes/ardrone_tracker.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python +import roslib +roslib.load_manifest('falkor_ardrone') + +import sys +import rospy +import cv +import cv2 +import math + +from std_msgs.msg import String +from sensor_msgs.msg import Image +from geometry_msgs.msg import Point +from cv_bridge import CvBridge, CvBridgeError + +import numpy as np +import tracker + +class ardroneTracker: + + def __init__(self, tracker): + self.point_pub = rospy.Publisher("/ardrone_tracker/found_point", Point ) + + self.bridge = CvBridge() + self.image_sub = rospy.Subscriber("/ardrone/front/image_raw",Image,self.callback) + self.image_pub = rospy.Publisher( "/ardrone_tracker/image", Image ) + + self.tracker = tracker + + def callback(self,data): + try: + cv_image = self.bridge.imgmsg_to_cv(data, "bgr8") + except CvBridgeError, e: + print e + + numpy_image = np.asarray( cv_image ) + trackData = self.tracker.track( numpy_image ) + if trackData: + x,y,z = trackData[0],trackData[1],trackData[2] + + point = Point( x,y,z ) + self.point_pub.publish( point ) + else: + self.point_pub.publish( Point( 0, 0, -1 ) ) + + try: + vis = self.tracker.get_vis() + image_message = self.bridge.cv_to_imgmsg(cv2.cv.fromarray( vis ), encoding="passthrough") + self.image_pub.publish( image_message ) + except CvBridgeError, e: + print e + +def main(): + rospy.init_node( 'ardrone_tracker' ) +# ardroneTracker(tracker.LkTracker() ) +# ardroneTracker(tracker.DummyTracker() ) +# '/usr/local/share/OpenCV/haarcascades/haarcascade_frontalface_alt.xml' + + cascade_file = rospy.get_param( '~cascadefile', 'NULL' ) + if cascade_file == 'NULL': + print "Must set cascadefile parameter!" + exit + + ardroneTracker( tracker.CascadeTracker( cascade_file ) ) + try: + rospy.spin() + except KeyboardInterrupt: + print "Shutting down" + cv.DestroyAllWindows() + +if __name__ == '__main__': +# import cProfile +# cProfile.run('main()') + main() diff --git a/nodes/pid.py b/nodes/pid.py new file mode 100644 index 0000000..a5974fa --- /dev/null +++ b/nodes/pid.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python + +class Pid2: + def __init__( self, gain_p = 0.0, gain_i = 0.0, gain_d = 0.0, + time_constant = 0.0, limit = -1.0 ): + self.gain_p = gain_p + self.gain_i = gain_i + self.gain_d = gain_d + self.time_constant = time_constant + self.limit = limit + + self.input = 0.0 + self.dinput = 0.0 + self.output = 0.0 + self.p = 0.0 + self.i = 0.0 + self.d = 0.0 + + def update( self, new_input, x, dx, dt ): + if self.limit > 0.0 and abs( new_input ) > self.limit: + if new_input < 0: + new_input = - self.limit + else: + new_input = self.limit + + if dt + self.time_constant > 0.0: + self.dinput = ( new_input - self.input ) / ( dt + self.time_constant ) + self.input = ( dt * new_input + self.time_constant * self.input ) / ( dt + self.time_constant ) + + self.p = self.input - x + self.d = self.dinput - dx + self.i = self.i + dt * self.p + +# print self.p,self.d,self.i + + self.output = self.gain_p * self.p + self.gain_d * self.d + self.gain_i * self.i + return self.output + + def reset( self ): + self.input = self.dinput = 0.0 + self.p = self.i = self.d = 0.0 + + +class Pid: + def __init__( self, Kp = 0.0, Ti = 0.0, Td = 0.0, outputLimit = None ): + self.Kp = Kp + self.Ti = Ti + self.Td = Td + + self.prev_error = 0.0 + self.integral = 0.0 + + self.outputLimit = outputLimit + self.setPointMin = self.setPointMax = 0.0 + + def get_output( self, pv, dt ): + if( pv < self.setPointMax and pv > self.setPointMin ): + error = 0 + else: + errorFromMax = self.setPointMax - pv + errorFromMin = self.setPointMin - pv + + if abs( errorFromMax ) > abs( errorFromMin ): + error = errorFromMin + else: + error = errorFromMax + + self.integral += error * dt + + if dt != 0: + derivative = ( error - self.prev_error ) / dt + else: + derivative = 0 + + self.prev_error = error + + if self.Ti == 0.0: + TiInv = 0 + else: + TiInv = 1/self.Ti + + output = self.Kp * ( error + TiInv * self.integral + self.Td * derivative ) + if self.outputLimit != None: + limitedOutput = min( max( output, -self.outputLimit ), + self.outputLimit ) + return limitedOutput + else: + return output diff --git a/nodes/tracker.py b/nodes/tracker.py new file mode 100644 index 0000000..680442e --- /dev/null +++ b/nodes/tracker.py @@ -0,0 +1,449 @@ +#!/usr/bin/env python +""" +Tracker classes must implement track( frame ) which returns a tuple (x,y,area) +if the tracker does not find an object then it will return None +""" + +import numpy as np +import math +import cv2 +import time + +lk_params = dict( winSize = ( 30, 30 ), + maxLevel = 2, + criteria = ( cv2.TERM_CRITERIA_EPS | + cv2.TERM_CRITERIA_COUNT, 10, 0.03 ) ) + +feature_params = dict( maxCorners = 50, + qualityLevel = 0.3, + minDistance = 10, + blockSize = 10 ) + +class DummyTracker: + def track( self, frame ): + (y,x,n) = frame.shape + cv2.imshow( 'Dummy', frame ) + return( 50, 50, 50 ) + +class LkTracker: + def __init__(self): + self.track_len = 10 + self.detect_interval = 5 + self.tracks = [] + self.frame_idx = 0 + self.frame = None + self.mouseDown = False + self.timers = {} + + self.userRect = None +# cv2.namedWindow( 'LKTracker', cv2.cv.CV_WINDOW_NORMAL ) +# cv2.cv.SetMouseCallback( 'LKTracker', self.on_mouse, None ) + + self.vis = None + + def on_mouse( self, event, x, y, flags, param ): + if self.frame is None: + return + + if event == cv2.cv.CV_EVENT_LBUTTONDOWN: + self.mouseDown = True + self.userRect = np.int32( ( ( x, y ), ( x, y ) ) ) + + elif event == cv2.cv.CV_EVENT_MOUSEMOVE and self.mouseDown == True: + xmin = min( self.userRect[0,0], self.userRect[1,0], x) + xmax = max( self.userRect[0,0], self.userRect[1,0], x) + ymin = min( self.userRect[0,1], self.userRect[1,1], y) + ymax = max( self.userRect[0,1], self.userRect[1,1], y) + self.userRect = np.int32( ( ( xmin, ymin ), ( xmax, ymax ) ) ) + + elif event == cv2.cv.CV_EVENT_LBUTTONUP: + self.mouseDown = False + self.pickFeatures() + self.initCamshift() + self.userRect = None + + def initCamshift(self): + x0,y0,x1,y1 = (self.userRect[0,0], + self.userRect[0,1], + self.userRect[1,0], + self.userRect[1,1]) + + hsv_roi = self.hsv[ y0:y1, x0:x1 ] + mask_roi = self.hsv_mask[ y0:y1, x0:x1 ] + + hist = cv2.calcHist( [hsv_roi], [0], mask_roi, [16], [0,180] ) + cv2.normalize( hist, hist, 0, 255, cv2.NORM_MINMAX ) + self.hist = hist.reshape(-1) + self.track_window = ( x0, y0, x1-x0, y1-y0 ) + self.showHist() + + def getCamShift(self): + prob = cv2.calcBackProject( [self.hsv], [0], self.hist, [0,180], 1 ) + prob &= self.hsv_mask + term_crit = ( cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 1 ) + track_box, self.track_window = cv2.CamShift( prob, self.track_window, + term_crit ) +# cv2.imshow( 'Back Projection', prob ) + return track_box + + def showHist(self): + bin_count = self.hist.shape[0] + bin_w = 24 + img = np.zeros((256, bin_count*bin_w, 3), np.uint8) + for i in xrange(bin_count): + h = int(self.hist[i]) + cv2.rectangle(img, (i*bin_w+2, 255), ((i+1)*bin_w-2, 255-h), (int(180.0*i/bin_count), 255, 255), -1) + img = cv2.cvtColor(img, cv2.COLOR_HSV2BGR) +# cv2.imshow('hist', img) + + def pickFeatures(self): + mask = np.zeros_like( self.frame_gray ) +# import pdb; pdb.set_trace() + + cv2.rectangle( mask, tuple( self.userRect[0] ), tuple( self.userRect[1] ), 255, -1 ) +# cv2.imshow( 'userMask', mask ) + + start = cv2.getTickCount() / cv2.getTickFrequency() + p = cv2.goodFeaturesToTrack( self.frame_gray, mask = mask, **feature_params ) + self.timers['GoodFeatures'] = cv2.getTickCount() / cv2.getTickFrequency() - start + + if p is not None: + self.tracks = [ [ (x,y) ] for x, y in np.float32(p).reshape(-1, 2) ] + + def equalizeHist(self, frame): + splits = cv2.split( frame ) + equalized = map( cv2.equalizeHist, splits ) + return cv2.merge( equalized ) + + def initializeFrame( self, frame ): + self.frame = cv2.pyrDown( frame ) + self.hsv = cv2.cvtColor( self.frame, cv2.COLOR_BGR2HSV ) + self.hsv = self.equalizeHist( self.hsv ) + self.hsv_mask = cv2.inRange( self.hsv, np.array((0.0, 33.0, 33.0)), + np.array((180.0,254.,254.)) ) + + # dilate + self.hsv_mask = cv2.dilate( self.hsv_mask, None, iterations = 5 ) + +# self.frame_gray = cv2.cvtColor(self.frame, cv2.COLOR_BGR2GRAY) + hsv_split = cv2.split( self.hsv ) + self.frame_gray = hsv_split[0] +# cv2.imshow( 'sat', hsv_split[1] ) +# cv2.imshow( 'vol', hsv_split[2] ) + + cv2.imshow( 'gray', self.frame_gray ) +# cv2.imshow( 'hsv_mask', self.hsv_mask ) + + def filterOutliers( self, deviations ): + pts = np.int32( [ tr[-1] for tr in self.tracks ] ) + if pts.size < 10: # 5 pts (10 bc x/y) + return + + x_median = np.median( pts[:,0] ) + y_median = np.median( pts[:,1] ) + + distances = ( np.square(pts[:,0] - x_median) + + np.square(pts[:,1] - y_median) ) + distance_median = np.median( distances ) + + self.tracks = [ tr for (i,tr) in enumerate( self.tracks ) + if distances[i] < distance_median * + np.square( deviations ) ] + + def runOpticalFlow( self ): + if len(self.tracks) > 0: + img0, img1 = self.prev_gray, self.frame_gray + p0 = np.float32([tr[-1] for tr in self.tracks]).reshape(-1, 1, 2 ) + start = cv2.getTickCount() / cv2.getTickFrequency() + p1, st, err = cv2.calcOpticalFlowPyrLK( img0, img1, p0, None, **lk_params ) + p0r, st, err = cv2.calcOpticalFlowPyrLK( img1, img0, p1, None, **lk_params ) + self.timers['LKTrack'] = cv2.getTickCount() / cv2.getTickFrequency() - start + + d = abs( p0-p0r ).reshape(-1, 2).max(-1) + good = d < 1 + new_tracks = [] + for tr, (x, y), good_flag in zip( self.tracks, p1.reshape(-1,2), good ): + if not good_flag: + continue + tr.append((x, y)) + if len(tr) > self.track_len: + del tr[0] + + new_tracks.append(tr) + + self.tracks = new_tracks + + def reDetect( self, use_camshift = True, expand_pixels = 1 ): + if self.frame_idx % self.detect_interval == 0 and len(self.tracks) > 0: + + if use_camshift: + mask_camshift = np.zeros_like( self.frame_gray ) + ellipse_camshift = self.getCamShift() + cv2.ellipse( mask_camshift, ellipse_camshift, 255, -1 ) + + mask_tracked = np.zeros_like( self.frame_gray ) + pts = np.int32( [ [tr[-1]] for tr in self.tracks ] ) + + if len(pts) > 5: + ellipse_tracked = cv2.fitEllipse( pts ) + else: + ellipse_tracked = cv2.minAreaRect( pts ) + + boundingRect = cv2.boundingRect( pts ) + # x0,y0,x1,y1 = ( boundingRect[0] - int(boundingRect[2] * 0.05), + # boundingRect[1] - int(boundingRect[3] * 0.05), + # boundingRect[0] + int(boundingRect[2] * 1.05), + # boundingRect[1] + int(boundingRect[3] * 1.05)) + + + x0,y0,x1,y1 = ( boundingRect[0] - expand_pixels, + boundingRect[1] - expand_pixels, + boundingRect[0] + boundingRect[2] + expand_pixels, + boundingRect[1] + boundingRect[3] + expand_pixels ) + + if len( pts ) > 2: + cv2.rectangle( mask_tracked, (x0,y0), (x1,y1), 255, -1 ) + + cv2.ellipse( mask_tracked, ellipse_tracked, 255, -1 ) + +# cv2.imshow( 'mask_camshift', mask_camshift ) +# cv2.imshow( 'mask_tracked', mask_tracked ) + + if use_camshift: + mask = mask_camshift & mask_tracked + else: + mask = mask_tracked + + for x, y in [np.int32(tr[-1]) for tr in self.tracks]: + cv2.circle(mask, (x,y), 5, 0, -1 ) + +# cv2.imshow( 'mask', mask ) + start = cv2.getTickCount() / cv2.getTickFrequency() + p = cv2.goodFeaturesToTrack( self.frame_gray, mask = mask, **feature_params ) + self.timers['GoodFeatures'] = cv2.getTickCount() / cv2.getTickFrequency() - start + + if p is not None: + for x, y in np.float32(p).reshape(-1, 2): + self.tracks.append([(x,y)]) + + def drawAndGetTrack( self ): + vis = self.frame.copy() + + if len(self.tracks) > 0: + cv2.polylines( vis, [ np.int32(tr) for tr in self.tracks ], False, ( 0, 255, 0 ) ) + pts = np.int32( [ [tr[-1]] for tr in self.tracks ] ) + + for [(x,y)] in pts: + cv2.circle( vis, (x, y), 2, (0, 255, 0), -1 ) + + if len( pts ) > 2: + boundingRect = cv2.boundingRect( pts ) + x0,y0,x1,y1 = ( boundingRect[0], boundingRect[1], + boundingRect[0] + boundingRect[2], + boundingRect[1] + boundingRect[3] ) + area = boundingRect[3] * boundingRect[2] + cx,cy = ( boundingRect[0] + boundingRect[2]/2, + boundingRect[1] + boundingRect[3]/2 ) + + cv2.rectangle( vis, (x0,y0), (x1,y1), (0,255,255), 3, 8, 0 ) + + cv2.circle( vis, (cx,cy), 5, (0,255,255), -1, 8, 0 ) + + if self.userRect != None: + cv2.rectangle( vis, tuple( self.userRect[0] ), tuple( self.userRect[1] ), ( 255, 255, 255 ) ) + + i = 0 + total = 0 + for ( timer, time ) in self.timers.items(): + i += 1 + total += time + cv2.putText( vis, '%s: %.1f ms' % ( timer, time * 1e3 ), + ( 20, i*20 ), cv2.FONT_HERSHEY_PLAIN, 1.0, + (255, 255, 255) ) + cv2.putText( vis, '%s: %.1f ms' % ( timer, time * 1e3 ), + ( 20, i*20 ), cv2.FONT_HERSHEY_PLAIN, 1.0, + ( 0, 0, 0), thickness = 2 ) + + + i += 1 + cv2.putText( vis, 'Total: %.1f ms' % ( total * 1e3 ), + ( 20, i*20 ), cv2.FONT_HERSHEY_PLAIN, 1.0, + (255, 255, 255) ) + cv2.putText( vis, 'Total: %.1f ms' % ( total * 1e3 ), + ( 20, i*20 ), cv2.FONT_HERSHEY_PLAIN, 1.0, + ( 0, 0, 0), thickness = 2 ) + + self.vis = vis.copy() + +# cv2.imshow( 'LKTracker', self.vis ) +# cv2.waitKey( 1 ) + + if len( self.tracks ) > 2: + imageSize = self.frame.shape + imageArea = imageSize[0]*imageSize[1] + + xRel = cx*100/imageSize[1] + yRel = cy*100/imageSize[0] + areaRel = math.sqrt( float( area ) / float( imageArea ) ) * 100 + return xRel,yRel,areaRel,cx,cy + else: + return None + + def get_vis( self ): + return self.vis + + def track(self, frame): + self.initializeFrame( frame ) + self.runOpticalFlow() + self.filterOutliers( 3 ) + + self.reDetect() + trackData = self.drawAndGetTrack() + + self.frame_idx += 1 + self.prev_gray = self.frame_gray + + return trackData + +class CascadeTracker( LkTracker ): + def __init__( self, cascadeFn ): + self.cascade = cv2.CascadeClassifier( cascadeFn ) + self.trackData = None + self.cascade_interval = 30 + + LkTracker.__init__( self ) + + def initCamshift( self ): + pass + + def track( self, frame ): + self.frame = frame # cv2.pyrDown( frame ) + self.frame_gray = cv2.cvtColor( self.frame, cv2.COLOR_BGR2GRAY ) + self.frame_gray = cv2.equalizeHist( self.frame_gray ) + + # do optical flow if we have it + self.runOpticalFlow() + self.trackData = self.drawAndGetTrack() + + # now look for the object every 5th frame + detectedObject = self.detectObject() + + # if we don't have it redetect optical flow + if detectedObject == None: + self.filterOutliers( 4 ) + self.reDetect( False, 0 ) + + # If we do have it, then re-select optical flow features + else: +# import pdb; pdb.set_trace() + self.userRect = [ [ detectedObject[0], + detectedObject[1] ], + [ detectedObject[0] + detectedObject[2], + detectedObject[1] + detectedObject[3] ] ] + self.pickFeatures() + + # Add points at the corners + corners = [ [ ( detectedObject[0], detectedObject[1] ) ], + [ ( detectedObject[0] + detectedObject[2], + detectedObject[1] ) ], + [ ( detectedObject[0], + detectedObject[1] + detectedObject[3] ) ], + [ ( detectedObject[0] + detectedObject[2], + detectedObject[1] + detectedObject[3] ) ] + ] + + self.tracks.extend( corners ) + + self.userRect = None + self.trackData = self.drawAndGetTrack() + + # imageSize = frame.shape + # imageArea = imageSize[0]*imageSize[1] + + # area = detectedObject[2] * detectedObject[3] + # cx = detectedObject[0] + detectedObject[2]/2 + # cy = detectedObject[1] + detectedObject[3]/2 + + # xRel = cx*100/imageSize[1] + # yRel = cy*100/imageSize[0] + # areaRel = math.sqrt( float( area ) / float( imageArea ) ) * 100 + # trackData = xRel,yRel,areaRel + + self.frame_idx += 1 + self.prev_gray = self.frame_gray + + return self.trackData + + def detectObject( self ): + # don't do a cascade every frame + if self.frame_idx % self.cascade_interval != 0: + return None + + start = cv2.getTickCount() / cv2.getTickFrequency() + objects = self.cascade.detectMultiScale( self.frame_gray, + scaleFactor=1.3, + minNeighbors=4, + minSize=(15, 15), + flags = cv2.cv.CV_HAAR_SCALE_IMAGE) + self.timers['Cascade'] = cv2.getTickCount() / cv2.getTickFrequency() - start + + vis = self.frame_gray.copy() + + if len( objects ) > 0: + for i in objects: + cv2.rectangle( vis, + ( int(i[0]), int(i[1]) ), + ( int(i[0]+i[2]), int(i[1]+i[3]) ), + cv2.cv.CV_RGB(0,255,0), 3, 8, 0) + +# cv2.imshow( "objects", vis ) + + # if we have only one object return that + # or if we have no trackData, return the first object + # if we have trackData and multiple objects, return the object + # closest to the trackData + if len( objects ) > 0: + if len( objects ) == 1 or not self.trackData: + return objects[0] + else: + distances_squared = ( np.square(objects[:,0]+objects[:,2]/2 - + self.trackData[3]) + + np.square(objects[:,1]+objects[:,2]/2 - + self.trackData[4]) ) + min_index = np.argmin( distances_squared ) + return objects[min_index] + else: + return None + +def main(): + import sys + try: video_src = sys.argv[1] + except: video_src = 0 + + tracker = CascadeTracker( + '/home/sameer/ros/falkor_ardrone/cascade/haarcascade_falkorlogopaper.xml' ) + cam = cv2.VideoCapture( video_src ) + + while True: + ret, frame = cam.read() + if ret: + trackData = tracker.track( frame ) + + if trackData: + print trackData + + ch = 0xFF & cv2.waitKey(1) + if ch == 27: + break + +if __name__ == '__main__': + main() + + + + + + + + +