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 as np
00006 import thread
00007 from sensor_msgs.msg import Image
00008 from geometry_msgs.msg import *
00009 from jsk_recognition_msgs.msg import Rect
00010 import cv2
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_recognition_msgs/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 = cv2.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,arr.shape[1]),
00043            min(status[2+idx][1]+ex_size[1]/2,arr.shape[0]))
00044     mask = np.ones((arr.shape[0], arr.shape[1]), dtype=np.uint8) * 255
00045     mask[pt1[0]:pt2[0], pt1[1]:pt2[1]] = 0
00046     status2 = cv2.minMaxLoc(arr, mask)
00047     return (status[0+idx],status2[0+idx],status[2+idx],status2[2+idx])
00048 
00049 def MinLock2nd(arr,ex_size):
00050     return _MinMaxLock2nd(arr,ex_size,True)
00051 def MaxLock2nd(arr,ex_size):
00052     return _MinMaxLock2nd(arr,ex_size,False)
00053 
00054 
00055 class MepConverter:
00056     bridge = CvBridge()
00057 
00058     def __init__(self):
00059         self.lockobj = thread.allocate_lock()
00060         # variables
00061         self.reference_image_msg = None
00062         self.templates = {} # frame_id : ref_image,ser_frame,ser_rect
00063         self.clock = [rospy.Time.now()]
00064 
00065         # publisher
00066         self.reference_pub = rospy.Publisher(
00067             "current_template", Image, queue_size=1)
00068         self.result_pub = rospy.Publisher(
00069             "result", TransformStamped, queue_size=1)
00070         self.debug_pub = rospy.Publisher("debug_image", Image, queue_size=1)
00071 
00072         # subscriber
00073         self.reference_image_sub = rospy.Subscriber(rospy.resolve_name("reference"),
00074                                                     Image,self.ref_image_callback,queue_size=1)
00075         self.search_image_sub = rospy.Subscriber(rospy.resolve_name("search"),
00076                                                   Image,self.ser_image_callback,queue_size=1)
00077         self.reference_point_sub = rospy.Subscriber(
00078             rospy.resolve_name("set_reference_point"),
00079             PointStamped,self.set_reference_point_callback,queue_size=1)
00080         self.search_sub = rospy.Subscriber(rospy.resolve_name("set_search_rect"),
00081                                            Rect,self.set_search_callback,queue_size=1)
00082 
00083         self.server = DynamicReconfigureServer(ConfigType, self.reconfigure)
00084 
00085     # general template modifier function
00086     def set_template (self, ref_id='', ref_image=None, ref_rect=None,
00087                       ser_frame=None, ser_rect=None, color=None, method=None):
00088         if not self.templates.has_key(ref_id):
00089             self.templates[ref_id] = {'ref_orig':None,'ref_point':None,'ref_image':None,
00090                                       'ser_frame':None, 'ser_rect':None,
00091                                       'color':'mono8', 'method':cv2.TM_SQDIFF_NORMED} # normalized SSD
00092             color_changed = True
00093         else:
00094             color_changed = (self.templates[ref_id]['color'] != color) and (color != None)
00095 
00096         something_changed = (ref_image is not None) or \
00097                             (ref_rect is not None) or color_changed
00098 
00099         # parameters for template
00100         if ref_image is not None:
00101             self.templates[ref_id]['ref_orig'] = ref_image.copy()
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 is not None and ref_image.ndim == 2:
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 is not None) and (ref_image is not None) and \
00114            something_changed:
00115             ref_rect = (min(max(0,ref_rect[0]),ref_image.shape[1]-ref_rect[2]),
00116                         min(max(0,ref_rect[1]),ref_image.shape[0]-ref_rect[3]),
00117                         ref_rect[2],ref_rect[3])
00118             ref_image_rect = ref_image[
00119                     ref_rect[1]:ref_rect[1] + ref_rect[3],
00120                     ref_rect[0]:ref_rect[0] + ref_rect[2]].copy()
00121             if color == 'bgr8' or (color == 'mono8' and ref_image.ndim == 2):
00122                 template_image = ref_image_rect
00123             elif color == 'mono8':
00124                 template_image = cv2.cvtColor(
00125                     ref_image_rect, cv2.COLOR_BGR2GRAY)
00126             elif color == 'hsv8':
00127                 template_image = cv2.cvtColor(
00128                     ref_image_rect, cv2.COLOR_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 = self.bridge.imgmsg_to_cv2(
00150                 self.reference_image_msg, "bgr8")
00151             self.set_template('',ref_image=cv_image, ref_rect=rect)
00152         except CvBridgeError, e:
00153             print e
00154 
00155     #
00156     # Callback Functions (lock require)
00157     #
00158 
00159     # simple point tracking callback
00160     def set_reference_point_callback (self, msg): # PointStamped
00161         self.lockobj.acquire()
00162         pt = (int(msg.point.x),int(msg.point.y))
00163         rect = (pt[0]-self.default_template_size[0]/2,
00164                 pt[1]-self.default_template_size[1]/2,
00165                 self.default_template_size[0], self.default_template_size[1])
00166         self.set_reference(rect)
00167         print rect
00168         search_rect = (pt[0]-self.default_search_size[0]/2,
00169                        pt[1]-self.default_search_size[1]/2,
00170                        self.default_search_size[0],self.default_search_size[1])
00171         self.set_template('',ser_frame=None, ser_rect=search_rect)
00172         self.lockobj.release()
00173 
00174     # store the latest image
00175     def set_search_callback (self, msg):
00176         self.lockobj.acquire()
00177         self.set_template(ser_rect=(msg.x,msg.y,msg.width,msg.height))
00178         self.lockobj.release()
00179 
00180     # store the latest image
00181     def ref_image_callback (self, msg):
00182         self.lockobj.acquire()
00183         self.reference_image_msg = msg
00184         self.lockobj.release()
00185 
00186     # reconfigure
00187     def reconfigure(self,config,level):
00188         self.lockobj.acquire()
00189         # param
00190         self.default_template_size = (config['default_template_width'],
00191                                       config['default_template_height'])
00192         self.default_search_size = (config['default_search_width'],
00193                                     config['default_search_height'])
00194         self.show_debug_image = config['show_debug_image']
00195         self.auto_search_area = config['auto_search_area']
00196 
00197         if config['current_template_id'] in self.templates.keys():
00198             # set template configuration
00199             template = self.templates[config['current_template_id']]
00200             method_enum = [
00201                 cv2.TM_SQDIFF, cv2.TM_SQDIFF_NORMED, cv2.TM_CCORR,
00202                 cv2.TM_CCORR_NORMED, cv2.TM_CCOEFF, cv2.TM_CCOEFF_NORMED]
00203             if (template['ref_point'] is not None) and \
00204                (template['ref_image'] is not None):
00205                 ref_pt = template['ref_point']
00206                 ref_size = template['ref_image'].shape
00207                 ref_rect = (ref_pt[0]-ref_size[0]/2,ref_pt[1]-ref_size[1]/2,ref_size[0],ref_size[1])
00208             self.set_template(ref_id=config['current_template_id'],
00209                               color=config['template_color_space'],
00210                               ref_rect=ref_rect,
00211                               method=method_enum[config['match_method']])
00212             # return current configuration
00213             config['template_color_method'] = template['color']
00214         self.lockobj.release()
00215         return config
00216 
00217     def ser_image_callback (self, msg):
00218         # initialize debug image
00219         if self.show_debug_image:
00220             debug_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
00221 
00222         self.lockobj.acquire()
00223 
00224         # match for each template
00225         for template_id in self.templates.keys():
00226             reference_image = self.templates[template_id]['ref_image']
00227             search_rect = self.templates[template_id]['ser_rect']
00228             matchmethod = self.templates[template_id]['method']
00229             color_space = self.templates[template_id]['color']
00230             if reference_image is None:
00231                 continue
00232             if search_rect == None: continue
00233             if (self.templates[template_id]['ser_frame'] != None and
00234                 self.templates[template_id]['ser_frame'] != msg.header.frame_id): continue
00235 
00236             # search rect &= image size
00237             search_rect = (max(0,search_rect[0]),
00238                            max(0,search_rect[1]),
00239                            min(msg.width,search_rect[0]+search_rect[2])- max(0,search_rect[0]),
00240                            min(msg.height,search_rect[1]+search_rect[3])- max(0,search_rect[1]))
00241 
00242             result = TransformStamped(header=msg.header)
00243 
00244             try:
00245                 if color_space == 'mono8':
00246                     search_image = self.bridge.imgmsg_to_cv2(msg, 'mono8')
00247                 else:
00248                     search_image = self.bridge.imgmsg_to_cv2(msg, 'bgr8')
00249                     if color_space != 'bgr8':
00250                         search_image = cv2.cvtColor(
00251                             search_image, cv2.COLOR_BGR2HSV)
00252 
00253                 image_size = search_image.shape
00254                 reference_size = (
00255                     reference_image.shape[1], reference_image.shape[0])
00256 
00257                 # search size < reference size
00258                 if (search_rect[2] < reference_size[0]) or (search_rect[3] < reference_size[1]):
00259                     continue
00260 
00261                 results = cv2.matchTemplate(
00262                     search_image[
00263                         search_rect[1]:search_rect[1] + search_rect[3],
00264                         search_rect[0]:search_rect[0] + search_rect[2]],
00265                     reference_image, matchmethod)
00266 
00267                 if matchmethod in [cv2.TM_SQDIFF, cv2.TM_SQDIFF_NORMED]:
00268                     status = MinLock2nd(results,reference_size) # minimum for SSD
00269                 else:
00270                     status = MaxLock2nd(results,reference_size) # maximum for others
00271                     status = (1 - status[0], 1 - status[1], status[2], status[3])
00272 
00273                 result_pt = (status[2][0]+search_rect[0]+reference_size[0]/2,
00274                              status[2][1]+search_rect[1]+reference_size[1]/2)
00275 
00276                 # publish center position as result
00277                 result.child_frame_id = template_id
00278                 result.transform.translation.x = result_pt[0]
00279                 result.transform.translation.y = result_pt[1]
00280                 result.transform.rotation.w = 1 # not rotate, temporary
00281                 self.result_pub.publish(result)
00282 
00283                 if reference_image.ndim == 2 and \
00284                    reference_image.dtype == np.uint8:
00285                     ref_msg = self.bridge.cv2_to_imgmsg(
00286                         reference_image, "mono8")
00287                 else:
00288                     ref_msg = self.bridge.cv2_to_imgmsg(
00289                         reference_image, "bgr8")
00290                 self.reference_pub.publish(ref_msg)
00291 
00292                 # self feedback
00293                 if self.auto_search_area:
00294                     val_x = result_pt[0]-search_rect[2]/2
00295                     val_y = result_pt[1]-search_rect[3]/2
00296                     ser_scale = max(2,5+np.log(status[0])) # ???
00297                     new_ser_rect = (
00298                         min(max(val_x,0),image_size[0]-search_rect[2]),
00299                         min(max(val_y,0),image_size[1]-search_rect[3]),
00300                         reference_size[0]*ser_scale,reference_size[1]*ser_scale)
00301                     self.set_template(ser_rect=new_ser_rect)
00302 
00303                 # draw on debug image
00304                 if self.show_debug_image:
00305                     cv2.rectangle(debug_image,
00306                                   (result_pt[0] - reference_size[0]/2,
00307                                    result_pt[1] - reference_size[1]/2),
00308                                   (result_pt[0] + reference_size[0]/2,
00309                                    result_pt[1] + reference_size[1]/2),
00310                                   color=(0, 0, 255))  # red
00311                     cv2.rectangle(debug_image,
00312                                   (search_rect[0], search_rect[1]),
00313                                   (search_rect[0] + search_rect[2],
00314                                    search_rect[1] + search_rect[3]),
00315                                   color=(0, 255, 0))
00316                     cv2.putText(
00317                         debug_image, str(status[0]),
00318                         (search_rect[0], search_rect[1] + search_rect[3]),
00319                         fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1,
00320                         color=(0, 255, 0), thickness=2,
00321                         lineType=16)  # 16 means cv2.CV_AA
00322 
00323                     if reference_image.ndim == 2:
00324                         reference_color_image = cv2.cvtColor(
00325                             reference_image, cv2.COLOR_GRAY2BGR)
00326                         debug_image[
00327                             :reference_size[1], :reference_size[0], :] = \
00328                             reference_color_image.copy()
00329                     else:
00330                         debug_image[
00331                             :reference_size[1], :reference_size[0], :] = \
00332                             reference_image.copy()
00333 
00334             except CvBridgeError, e:
00335                 print e
00336 
00337         self.lockobj.release()
00338 
00339         if self.show_debug_image:
00340             # calc and print fps (30frame avg)
00341             self.clock += [rospy.Time.now()]
00342             if 30 <= len(self.clock):
00343                 fps = 30.0 / (self.clock[29]-self.clock[0]).to_sec()
00344                 cv2.putText(
00345                     debug_image, '%.1f fps' % fps, (msg.width - 150, 40),
00346                     fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1,
00347                     color=(0, 0, 255),  # red
00348                     thickness=2, lineType=16)  # 16 means cv2.CV_AA
00349                 self.clock = self.clock[-30:]
00350             # publish debug image
00351             debug_msg = self.bridge.cv2_to_imgmsg(debug_image, "bgr8")
00352             self.debug_pub.publish(debug_msg)
00353 
00354         return
00355 
00356 if __name__=='__main__':
00357     rospy.init_node('match_template')
00358 
00359     obj = MepConverter()
00360 
00361     ## reconfigure
00362     rospy.loginfo("reference image %s, size %s", obj.reference_image_sub.resolved_name, obj.default_template_size)
00363     rospy.loginfo("   search_image %s, size %s", obj.search_image_sub.resolved_name, obj.default_search_size)
00364 
00365     rospy.spin()


jsk_perception
Author(s): Manabu Saito, Ryohei Ueda
autogenerated on Tue Jul 2 2019 19:41:07