matchtemplate.py
Go to the documentation of this file.
00001 #!/usr/bin/env python
00002 
00003 import roslib; roslib.load_manifest('jsk_perception')
00004 import rospy
00005 import numpy
00006 import thread
00007 from sensor_msgs.msg import Image
00008 from geometry_msgs.msg import *
00009 from jsk_perception.msg import *
00010 import cv
00011 from cv_bridge import CvBridge, CvBridgeError
00012 
00013 from dynamic_reconfigure.server import Server as DynamicReconfigureServer
00014 from dynamic_reconfigure.msg import SensorLevels
00015 from jsk_perception.cfg import matchtemplateConfig as ConfigType
00016 
00017 # you can set below options for each templates
00018 #   color space: mono8,bgr8,bgra8,hsv,hsva (OpenCV expects bgr)
00019 #   template size: tuple (x,y,width,height)
00020 #   search area:   tuple (x,y,width,height)
00021 #   search topic: (TODO)
00022 #   match method: 6type
00023 
00024 # publish topic:
00025 #   ~/current_template    [sensor_msgs/Image]
00026 #   ~/result              [geometry_msgs/TransformStamped]
00027 #   ~/debug_image         [sensor_msgs/Image]
00028 
00029 # subscribe topic:
00030 #   ~/reference           [sensor_msgs/Image]
00031 #   ~/search              [sensor_msgs/Image]
00032 #   ~/set_reference_point [geometry_msgs/PointStamped]
00033 #   ~/set_search_rect     [jsk_perception/Rect]
00034 
00035 # return (1st-val,1st-loc,2nd-val,2nd-loc)
00036 def _MinMaxLock2nd(arr,ex_size,is_min):
00037     if is_min: idx = 0
00038     else: idx = 1
00039     status = cv.MinMaxLoc(arr)
00040     pt1 = (max(status[2+idx][0]-ex_size[0]/2,0),
00041            max(status[2+idx][1]-ex_size[1]/2,0))
00042     pt2 = (min(status[2+idx][0]+ex_size[0]/2,cv.GetSize(arr)[0]),
00043            min(status[2+idx][1]+ex_size[1]/2,cv.GetSize(arr)[1]))
00044     mask = cv.CreateImage(cv.GetSize(arr), cv.IPL_DEPTH_8U, 1)
00045     cv.Set(mask,cv.Scalar(255))
00046     cv.SetImageROI(mask,(pt1[0],pt1[1],pt2[0]-pt1[0],pt2[1]-pt1[1]))
00047     cv.Set(mask,cv.Scalar(0))
00048     cv.ResetImageROI(mask)
00049     status2 = cv.MinMaxLoc(arr,mask)
00050     return (status[0+idx],status2[0+idx],status[2+idx],status2[2+idx])
00051 
00052 def MinLock2nd(arr,ex_size):
00053     return _MinMaxLock2nd(arr,ex_size,True)
00054 def MaxLock2nd(arr,ex_size):
00055     return _MinMaxLock2nd(arr,ex_size,False)
00056 
00057 
00058 class MepConverter:
00059     bridge = CvBridge()
00060 
00061     def __init__(self):
00062         self.lockobj = thread.allocate_lock()
00063         # variables
00064         self.reference_image_msg = None
00065         self.templates = {} # frame_id : ref_image,ser_frame,ser_rect
00066         self.clock = [rospy.Time.now()]
00067 
00068         # publisher
00069         self.reference_pub = rospy.Publisher("current_template",Image)
00070         self.result_pub = rospy.Publisher("result",TransformStamped)
00071         self.debug_pub = rospy.Publisher("debug_image",Image)
00072 
00073         # subscriber
00074         self.reference_image_sub = rospy.Subscriber(rospy.resolve_name("reference"),
00075                                                     Image,self.ref_image_callback,queue_size=1)
00076         self.search_image_sub = rospy.Subscriber(rospy.resolve_name("search"),
00077                                                   Image,self.ser_image_callback,queue_size=1)
00078         self.reference_point_sub = rospy.Subscriber(
00079             rospy.resolve_name("set_reference_point"),
00080             PointStamped,self.set_reference_point_callback,queue_size=1)
00081         self.search_sub = rospy.Subscriber(rospy.resolve_name("set_search_rect"),
00082                                            Rect,self.set_search_callback,queue_size=1)
00083 
00084         self.server = DynamicReconfigureServer(ConfigType, self.reconfigure)
00085 
00086     # general template modifier function
00087     def set_template (self, ref_id='', ref_image=None, ref_rect=None,
00088                       ser_frame=None, ser_rect=None, color=None, method=None):
00089         if not self.templates.has_key(ref_id):
00090             self.templates[ref_id] = {'ref_orig':None,'ref_point':None,'ref_image':None,
00091                                       'ser_frame':None, 'ser_rect':None,
00092                                       'color':'mono8', 'method':cv.CV_TM_SQDIFF_NORMED} # normalized SSD
00093             color_changed = True
00094         else:
00095             color_changed = (self.templates[ref_id]['color'] != color) and (color != None)
00096 
00097         something_changed = (ref_image != None) or (ref_rect != None) or color_changed
00098 
00099         # parameters for template
00100         if ref_image != None:
00101             self.templates[ref_id]['ref_orig'] = cv.CloneImage(ref_image)
00102         else:
00103             ref_image = self.templates[ref_id]['ref_orig'] # old image can be set here
00104 
00105         if color != None:
00106             if ref_image != None and ref_image.nChannels == 1:
00107                 color = 'mono8'
00108             self.templates[ref_id]['color'] = color
00109         else:
00110             color = self.templates[ref_id]['color']
00111 
00112         # copy template from reference original image
00113         if (ref_rect != None) and (ref_image != None) and something_changed:
00114             ref_rect = (min(max(0,ref_rect[0]),ref_image.width-ref_rect[2]),
00115                         min(max(0,ref_rect[1]),ref_image.height-ref_rect[3]),
00116                         ref_rect[2],ref_rect[3])
00117             template_image = cv.CreateImage((ref_rect[2], ref_rect[3]),
00118                                          ref_image.depth,
00119                                          {'mono8':1,'bgr8':3,'hsv8':3}[color])
00120             cv.SetImageROI(ref_image, ref_rect)
00121             if color == 'mono8' and ref_image.nChannels == 1:
00122                 cv.Copy(ref_image, template_image)
00123             elif color == 'mono8':
00124                 cv.CvtColor(ref_image,template_image,cv.CV_BGR2GRAY)
00125             elif color == 'bgr8':
00126                 cv.Copy(ref_image, template_image)
00127             elif color == 'hsv8':
00128                 cv.CvtColor(ref_image,template_image,cv.CV_BGR2HSV)
00129 
00130             self.templates[ref_id]['ref_point'] = (ref_rect[0]+ref_rect[2]/2,
00131                                                    ref_rect[1]+ref_rect[3]/2)
00132             self.templates[ref_id]['ref_image'] = template_image
00133             rospy.loginfo("set ref_image id=%s, rect=%s", ref_id, ref_rect);
00134 
00135         # matching parameters
00136         if ser_frame != None:
00137             self.templates[ref_id]['ser_frame'] = ser_frame
00138             rospy.loginfo("set ser_frame id=%s frame=%s", ref_id, ser_frame);
00139         if ser_rect != None:
00140             self.templates[ref_id]['ser_rect'] = ser_rect
00141             #rospy.loginfo("set ser_rect id=%s %s", ref_id, ser_rect);
00142         if method != None:
00143             self.templates[ref_id]['method'] = method
00144             rospy.loginfo("set method id=%s method=%s", ref_id, method);
00145 
00146     def set_reference (self, rect):
00147         if self.reference_image_msg == None: return
00148         try:
00149             cv_image = cv.GetImage(self.bridge.imgmsg_to_cv(self.reference_image_msg, "bgr8"))
00150             self.set_template('',ref_image=cv_image, ref_rect=rect)
00151         except CvBridgeError, e:
00152             print e
00153 
00154     #
00155     # Callback Functions (lock require)
00156     #
00157 
00158     # simple point tracking callback
00159     def set_reference_point_callback (self, msg): # PointStamped
00160         self.lockobj.acquire()
00161         pt = (int(msg.point.x),int(msg.point.y))
00162         rect = (pt[0]-self.default_template_size[0]/2,
00163                 pt[1]-self.default_template_size[1]/2,
00164                 self.default_template_size[0], self.default_template_size[1])
00165         self.set_reference(rect)
00166         print rect
00167         search_rect = (pt[0]-self.default_search_size[0]/2,
00168                        pt[1]-self.default_search_size[1]/2,
00169                        self.default_search_size[0],self.default_search_size[1])
00170         self.set_template('',ser_frame=None, ser_rect=search_rect)
00171         self.lockobj.release()
00172 
00173     # store the latest image
00174     def set_search_callback (self, msg):
00175         self.lockobj.acquire()
00176         self.set_template(ser_rect=(msg.x,msg.y,msg.width,msg.height))
00177         self.lockobj.release()
00178 
00179     # store the latest image
00180     def ref_image_callback (self, msg):
00181         self.lockobj.acquire()
00182         self.reference_image_msg = msg
00183         self.lockobj.release()
00184 
00185     # reconfigure
00186     def reconfigure(self,config,level):
00187         self.lockobj.acquire()
00188         # param
00189         self.default_template_size = (config['default_template_width'],
00190                                       config['default_template_height'])
00191         self.default_search_size = (config['default_search_width'],
00192                                     config['default_search_height'])
00193         self.show_debug_image = config['show_debug_image']
00194         self.auto_search_area = config['auto_search_area']
00195 
00196         if config['current_template_id'] in self.templates.keys():
00197             # set template configuration
00198             template = self.templates[config['current_template_id']]
00199             method_enum = [cv.CV_TM_SQDIFF, cv.CV_TM_SQDIFF_NORMED, cv.CV_TM_CCORR,
00200                            cv.CV_TM_CCORR_NORMED, cv.CV_TM_CCOEFF, cv.CV_TM_CCOEFF_NORMED]
00201             if (template['ref_point'] != None) and (template['ref_image'] != None):
00202                 ref_pt = template['ref_point']
00203                 ref_size = cv.GetSize(template['ref_image'])
00204                 ref_rect = (ref_pt[0]-ref_size[0]/2,ref_pt[1]-ref_size[1]/2,ref_size[0],ref_size[1])
00205             self.set_template(ref_id=config['current_template_id'],
00206                               color=config['template_color_space'],
00207                               ref_rect=ref_rect,
00208                               method=method_enum[config['match_method']])
00209             # return current configuration
00210             config['template_color_method'] = template['color']
00211         self.lockobj.release()
00212         return config
00213 
00214     def ser_image_callback (self, msg):
00215         # initialize debug image
00216         if self.show_debug_image:
00217             debug_image = cv.CloneImage(cv.GetImage(self.bridge.imgmsg_to_cv(msg, "bgr8")))
00218 
00219         self.lockobj.acquire()
00220 
00221         # match for each template
00222         for template_id in self.templates.keys():
00223             reference_image = self.templates[template_id]['ref_image']
00224             search_rect = self.templates[template_id]['ser_rect']
00225             matchmethod = self.templates[template_id]['method']
00226             color_space = self.templates[template_id]['color']
00227             if reference_image == None: continue
00228             if search_rect == None: continue
00229             if (self.templates[template_id]['ser_frame'] != None and
00230                 self.templates[template_id]['ser_frame'] != msg.header.frame_id): continue
00231 
00232             # search rect &= image size
00233             search_rect = (max(0,search_rect[0]),
00234                            max(0,search_rect[1]),
00235                            min(msg.width,search_rect[0]+search_rect[2])- max(0,search_rect[0]),
00236                            min(msg.height,search_rect[1]+search_rect[3])- max(0,search_rect[1]))
00237 
00238             result = TransformStamped(header=msg.header)
00239 
00240             try:
00241                 if color_space in ['mono8','bgr8']:
00242                     search_image = cv.GetImage(self.bridge.imgmsg_to_cv(msg, color_space))
00243                 else:
00244                     cv_image = cv.GetImage(self.bridge.imgmsg_to_cv(msg, 'bgr8'))
00245                     search_image = cv.CreateImage(cv.GetSize(cv_image),cv_image.depth,cv_image.nChannels)
00246                     cv.CvtColor(cv_image,search_image,cv.CV_BGR2HSV)
00247 
00248                 image_size = cv.GetSize(search_image)
00249                 reference_size = cv.GetSize(reference_image)
00250 
00251                 # search size < reference size
00252                 if (search_rect[2] < reference_size[0]) or (search_rect[3] < reference_size[1]):
00253                     continue
00254 
00255                 cv.SetImageROI(search_image,search_rect)
00256                 ressize = list(search_rect)[2:]
00257                 ressize[0] -= reference_size[0] - 1
00258                 ressize[1] -= reference_size[1] - 1
00259                 results = cv.CreateImage(ressize, cv.IPL_DEPTH_32F, 1)
00260                 cv.MatchTemplate(search_image, reference_image, results, matchmethod)
00261                 cv.ResetImageROI(search_image)
00262 
00263                 if matchmethod in [cv.CV_TM_SQDIFF,cv.CV_TM_SQDIFF_NORMED]:
00264                     status = MinLock2nd(results,reference_size) # minimum for SSD
00265                 else:
00266                     status = MaxLock2nd(results,reference_size) # maximum for others
00267                     status = (1 - status[0], 1 - status[1], status[2], status[3])
00268 
00269                 result_pt = (status[2][0]+search_rect[0]+reference_size[0]/2,
00270                              status[2][1]+search_rect[1]+reference_size[1]/2)
00271 
00272                 # publish center position as result
00273                 result.child_frame_id = template_id
00274                 result.transform.translation.x = result_pt[0]
00275                 result.transform.translation.y = result_pt[1]
00276                 result.transform.rotation.w = 1 # not rotate, temporary
00277                 self.result_pub.publish(result)
00278 
00279                 self.reference_pub.publish(self.bridge.cv_to_imgmsg(reference_image, "passthrough"))
00280 
00281                 # self feedback
00282                 if self.auto_search_area:
00283                     val_x = result_pt[0]-search_rect[2]/2
00284                     val_y = result_pt[1]-search_rect[3]/2
00285                     ser_scale = max(2,5+numpy.log(status[0])) # ???
00286                     new_ser_rect = (
00287                         min(max(val_x,0),image_size[0]-search_rect[2]),
00288                         min(max(val_y,0),image_size[1]-search_rect[3]),
00289                         reference_size[0]*ser_scale,reference_size[1]*ser_scale)
00290                     self.set_template(ser_rect=new_ser_rect)
00291 
00292                 # draw on debug image
00293                 if self.show_debug_image:
00294                     cv.Rectangle(debug_image,
00295                               (result_pt[0] - reference_size[0]/2,
00296                                result_pt[1] - reference_size[1]/2),
00297                               (result_pt[0] + reference_size[0]/2,
00298                                result_pt[1] + reference_size[1]/2),
00299                               cv.CV_RGB(255,0,0))
00300                     cv.Rectangle(debug_image,
00301                               (search_rect[0], search_rect[1]),
00302                               (search_rect[0] + search_rect[2],
00303                                search_rect[1] + search_rect[3]),
00304                               cv.CV_RGB(0,255,0))
00305                     cv.PutText(debug_image, str(status[0]), (search_rect[0], search_rect[1] + search_rect[3]), font, cv.CV_RGB(0,255,0))
00306 
00307                     cv.SetImageROI(debug_image,(0,0,reference_size[0],reference_size[1]))
00308                     if reference_image.nChannels == 1:
00309                         reference_color_image = cv.CreateImage(reference_size, debug_image.depth, debug_image.nChannels)
00310                         cv.CvtColor (reference_image, reference_color_image, cv.CV_GRAY2BGR) # TODO
00311                         cv.Copy(reference_color_image,debug_image)
00312                     else:
00313                         cv.Copy(reference_image,debug_image)
00314                     cv.ResetImageROI(debug_image)
00315 
00316             except CvBridgeError, e:
00317                 print e
00318 
00319         self.lockobj.release()
00320 
00321         if self.show_debug_image:
00322             # calc and print fps (30frame avg)
00323             self.clock += [rospy.Time.now()]
00324             if 30 <= len(self.clock):
00325                 fps = 30.0 / (self.clock[29]-self.clock[0]).to_sec()
00326                 cv.PutText(debug_image, '%.1f fps'%fps, (msg.width-150,40),font,cv.CV_RGB(255,0,0))
00327                 self.clock = self.clock[-30:]
00328             # publish debug image
00329             self.debug_pub.publish(self.bridge.cv_to_imgmsg(debug_image, "bgr8"))
00330 
00331         return
00332 
00333 if __name__=='__main__':
00334     rospy.init_node('match_template')
00335 
00336     obj = MepConverter()
00337 
00338     ## reconfigure
00339     rospy.loginfo("reference image %s, size %s", obj.reference_image_sub.resolved_name, obj.default_template_size)
00340     rospy.loginfo("   search_image %s, size %s", obj.search_image_sub.resolved_name, obj.default_search_size)
00341 
00342     # opencv
00343     font = cv.InitFont(cv.CV_FONT_HERSHEY_SIMPLEX, 1, 1, 0.0, 2, cv.CV_AA)
00344 
00345     rospy.spin()


jsk_perception
Author(s): Manabu Saito
autogenerated on Mon Oct 6 2014 01:16:59