$search
00001 #!/usr/bin/env python 00002 00003 """ camshift_node.py - Version 1.0 2011-04-19 00004 00005 Modification of the ROS OpenCV Camshift example using cv_bridge and publishing the ROI 00006 coordinates to the /roi topic. 00007 """ 00008 00009 import roslib 00010 roslib.load_manifest('pi_head_tracking_tutorial') 00011 import sys 00012 import rospy 00013 import cv 00014 from std_msgs.msg import String 00015 from sensor_msgs.msg import Image, RegionOfInterest, CameraInfo 00016 from cv_bridge import CvBridge, CvBridgeError 00017 00018 class CamShiftNode: 00019 def __init__(self): 00020 rospy.init_node('cam_shift_node') 00021 00022 self.ROI = rospy.Publisher("roi", RegionOfInterest) 00023 00024 """ Create the display window """ 00025 self.cv_window_name = "Camshift Tracker" 00026 cv.NamedWindow(self.cv_window_name, 0) 00027 00028 """ Create the cv_bridge object """ 00029 self.bridge = CvBridge() 00030 00031 """ Subscribe to the raw camera image topic """ 00032 self.image_sub = rospy.Subscriber("/camera/image_raw", Image, self.image_callback) 00033 00034 """ Set up a smaller window to display the CamShift histogram. """ 00035 cv.NamedWindow("Histogram", 0) 00036 cv.MoveWindow("Histogram", 700, 10) 00037 cv.SetMouseCallback(self.cv_window_name, self.on_mouse) 00038 00039 self.drag_start = None # Set to (x,y) when mouse starts dragtime 00040 self.track_window = None # Set to rect when the mouse drag finishes 00041 00042 self.hist = cv.CreateHist([180], cv.CV_HIST_ARRAY, [(0,180)], 1 ) 00043 self.backproject_mode = False 00044 00045 def image_callback(self, data): 00046 """ Convert the raw image to OpenCV format using the convert_image() helper function """ 00047 cv_image = self.convert_image(data) 00048 00049 """ Apply the CamShift algorithm using the do_camshift() helper function """ 00050 cv_image = self.do_camshift(cv_image) 00051 00052 """ Refresh the displayed image """ 00053 cv.ShowImage(self.cv_window_name, cv_image) 00054 00055 """ Toggle between the normal and back projected image if user hits the 'b' key """ 00056 c = cv.WaitKey(7) % 0x100 00057 if c == 27: 00058 return 00059 elif c == ord("b"): 00060 self.backproject_mode = not self.backproject_mode 00061 00062 def convert_image(self, ros_image): 00063 try: 00064 cv_image = self.bridge.imgmsg_to_cv(ros_image, "bgr8") 00065 return cv_image 00066 except CvBridgeError, e: 00067 print e 00068 00069 def do_camshift(self, cv_image): 00070 """ Get the image size """ 00071 image_size = cv.GetSize(cv_image) 00072 image_width = image_size[0] 00073 image_height = image_size[1] 00074 00075 """ Convert to HSV and keep the hue """ 00076 hsv = cv.CreateImage(image_size, 8, 3) 00077 cv.CvtColor(cv_image, hsv, cv.CV_BGR2HSV) 00078 self.hue = cv.CreateImage(image_size, 8, 1) 00079 cv.Split(hsv, self.hue, None, None, None) 00080 00081 """ Compute back projection """ 00082 backproject = cv.CreateImage(image_size, 8, 1) 00083 00084 """ Run the cam-shift algorithm """ 00085 cv.CalcArrBackProject( [self.hue], backproject, self.hist ) 00086 if self.track_window and is_rect_nonzero(self.track_window): 00087 crit = ( cv.CV_TERMCRIT_EPS | cv.CV_TERMCRIT_ITER, 10, 1) 00088 (iters, (area, value, rect), track_box) = cv.CamShift(backproject, self.track_window, crit) 00089 self.track_window = rect 00090 00091 """ If mouse is pressed, highlight the current selected rectangle 00092 and recompute the histogram """ 00093 00094 if self.drag_start and is_rect_nonzero(self.selection): 00095 sub = cv.GetSubRect(cv_image, self.selection) 00096 save = cv.CloneMat(sub) 00097 cv.ConvertScale(cv_image, cv_image, 0.5) 00098 cv.Copy(save, sub) 00099 x,y,w,h = self.selection 00100 cv.Rectangle(cv_image, (x,y), (x+w,y+h), (255,255,255)) 00101 00102 sel = cv.GetSubRect(self.hue, self.selection ) 00103 cv.CalcArrHist( [sel], self.hist, 0) 00104 (_, max_val, _, _) = cv.GetMinMaxHistValue(self.hist) 00105 if max_val != 0: 00106 cv.ConvertScale(self.hist.bins, self.hist.bins, 255. / max_val) 00107 elif self.track_window and is_rect_nonzero(self.track_window): 00108 cv.EllipseBox( cv_image, track_box, cv.CV_RGB(255,0,0), 3, cv.CV_AA, 0 ) 00109 00110 roi = RegionOfInterest() 00111 roi.x_offset = int(min(image_width, max(0, track_box[0][0] - track_box[1][0] / 2))) 00112 roi.y_offset = int(min(image_height, max(0, track_box[0][1] - track_box[1][1] / 2))) 00113 roi.width = int(track_box[1][0]) 00114 roi.height = int(track_box[1][1]) 00115 self.ROI.publish(roi) 00116 00117 cv.ShowImage("Histogram", self.hue_histogram_as_image(self.hist)) 00118 00119 if not self.backproject_mode: 00120 return cv_image 00121 else: 00122 return backproject 00123 00124 00125 def hue_histogram_as_image(self, hist): 00126 """ Returns a nice representation of a hue histogram """ 00127 00128 histimg_hsv = cv.CreateImage( (320,200), 8, 3) 00129 00130 mybins = cv.CloneMatND(hist.bins) 00131 cv.Log(mybins, mybins) 00132 (_, hi, _, _) = cv.MinMaxLoc(mybins) 00133 cv.ConvertScale(mybins, mybins, 255. / hi) 00134 00135 w,h = cv.GetSize(histimg_hsv) 00136 hdims = cv.GetDims(mybins)[0] 00137 for x in range(w): 00138 xh = (180 * x) / (w - 1) # hue sweeps from 0-180 across the image 00139 val = int(mybins[int(hdims * x / w)] * h / 255) 00140 cv.Rectangle( histimg_hsv, (x, 0), (x, h-val), (xh,255,64), -1) 00141 cv.Rectangle( histimg_hsv, (x, h-val), (x, h), (xh,255,255), -1) 00142 00143 histimg = cv.CreateImage( (320,200), 8, 3) 00144 cv.CvtColor(histimg_hsv, histimg, cv.CV_HSV2BGR) 00145 return histimg 00146 00147 def on_mouse(self, event, x, y, flags, param): 00148 if event == cv.CV_EVENT_LBUTTONDOWN: 00149 self.drag_start = (x, y) 00150 if event == cv.CV_EVENT_LBUTTONUP: 00151 self.drag_start = None 00152 self.track_window = self.selection 00153 if self.drag_start: 00154 xmin = min(x, self.drag_start[0]) 00155 ymin = min(y, self.drag_start[1]) 00156 xmax = max(x, self.drag_start[0]) 00157 ymax = max(y, self.drag_start[1]) 00158 self.selection = (xmin, ymin, xmax - xmin, ymax - ymin) 00159 00160 def is_rect_nonzero(r): 00161 (_,_,w,h) = r 00162 return (w > 0) and (h > 0) 00163 00164 def main(args): 00165 cs = CamShiftNode() 00166 try: 00167 rospy.spin() 00168 except KeyboardInterrupt: 00169 print "Shutting down vision node." 00170 cv.DestroyAllWindows() 00171 00172 if __name__ == '__main__': 00173 main(sys.argv)