test14_image_matching.py
Go to the documentation of this file.
00001 import roslib; roslib.load_manifest('hai_sandbox')
00002 import rospy
00003 import hrl_camera.ros_camera as rc
00004 import cv
00005 import hai_sandbox.features as fea
00006 import sys
00007 import threading as thr
00008 import scipy.spatial as sp
00009 import numpy as np
00010 import pdb
00011 
00012 class ShowImage(thr.Thread):
00013     def __init__(self, name):
00014         thr.Thread.__init__(self)
00015         self.image = None
00016         self.name = name
00017 
00018     def run(self):
00019         while not rospy.is_shutdown():
00020             if self.image != None:
00021                 cv.ShowImage(self.name, self.image)
00022                 cv.WaitKey(33)
00023 
00024 def concat_images(a, b):
00025     img_height = max(a.height, b.height)
00026     c = cv.CreateImage((a.width+b.width, img_height), a.depth, a.channels)
00027     a_area = cv.GetSubRect(c, (0,0, a.width, a.height))
00028     b_area = cv.GetSubRect(c, (a.width, 0, b.width, b.height))
00029     cv.Add(a, a_area, a_area)
00030     cv.Add(b, b_area, b_area)
00031     return c
00032 
00033 
00034 #class SURFMatcher:
00035 #    def __init__(self):
00036 #        self.model_images = {}
00037 #        self.model_fea = {}
00038 #
00039 #    def add_file(self, model_name, label): 
00040 #        model_img = cv.LoadImage(model_name)
00041 #        self.add_model(model_img, label)
00042 #
00043 #    def add_model(self, model_img, label):
00044 #        mgray = fea.grayscale(model_img)
00045 #        m_loc, m_desc = fea.surf(mgray)
00046 #        self.model_images[label] = model_img
00047 #        self.model_fea[label] = {'loc': m_loc, 'desc': m_desc}
00048 #
00049 #    def build_db(self):
00050 #        fea_l = []
00051 #        labels_l = []
00052 #        locs_l = []
00053 #        for k in self.model_fea:
00054 #            fea_l.append(np.array(self.model_fea[k]['desc']))
00055 #            locs_l.append(np.array(self.model_fea[k]['loc']))
00056 #            labels_l.append(np.array([k for i in range(len(self.model_fea[k]['desc']))]))
00057 #
00058 #        self.labels = np.row_stack(labels_l)
00059 #        self.locs = np.row_stack(locs_l)
00060 #        self.tree = sp.KDTree(np.row_stack(fea_l))
00061 #
00062 #    def match(self, desc, thres=.6):
00063 #        dists, idxs = self.tree.query(np.array(desc), 2)
00064 #        ratio = dists[0] / dists[1]
00065 #        if ratio < threshold:
00066 #            desc = self.tree.data[idxs[0]]
00067 #            loc = self.locs[idxs[0]]
00068 #            return desc, loc
00069 #        else:
00070 #            return None
00071 
00072 
00073 def match_images(model_img, cand_img, threshold=.8):
00074     #pdb.set_trace()
00075     mgray = fea.grayscale(model_img)
00076     cgray = fea.grayscale(cand_img)
00077     
00078     m_loc, m_desc = fea.surf(mgray)
00079     dirs = [direction for loc, lap, size, direction, hess in m_loc]
00080     print 'max min dirs', np.min(dirs), np.max(dirs)
00081 
00082     c_loc, c_desc = fea.surf(cgray)
00083     
00084     features_db = sp.KDTree(np.array(m_desc))
00085     matched = []
00086     for i, desc in enumerate(c_desc):
00087         dists, idxs = features_db.query(np.array(desc), 2)
00088         ratio = dists[0] / dists[1]
00089         #print "%d %.4f" % (i, ratio),
00090         if ratio < threshold:
00091             matched.append((i, idxs[0]))
00092             #print 'matched!', idxs[0]
00093         #else:
00094         #    print 'X|'
00095     
00096     c_loc_moved = []
00097     for loc, lap, size, d, hess in c_loc:
00098         x, y = loc
00099         nloc = (x + model_img.width, y)
00100         c_loc_moved.append((nloc, lap, size, d, hess))
00101     
00102     c_loc_matched, m_loc_matched = zip(*[[c_loc_moved[i], m_loc[j]] for i, j in matched])
00103     joint = concat_images(model_img, cand_img)
00104     
00105     joint_viz = joint
00106     #joint_viz = fea.draw_surf(joint, c_loc_moved, (255,0,0))
00107     #joint_viz = fea.draw_surf(joint_viz, c_loc_matched, (0,255,0))
00108     #joint_viz = fea.draw_surf(joint_viz, m_loc, (255,0,0))
00109     #joint_viz = fea.draw_surf(joint_viz, m_loc_matched, (0,255,0))
00110     for cloc, mloc in zip(c_loc_matched, m_loc_matched):
00111         cloc2d, _, _, _, _ = cloc
00112         mloc2d, _, _, _, _ = mloc
00113         cv.Line(joint_viz, cloc2d, mloc2d, (0,255,0), 1, cv.CV_AA)
00114     print '%d matches found' % len(matched)
00115     
00116     return joint_viz
00117 
00118 def test_thresholds():
00119     model_name = sys.argv[1]
00120     candidate = sys.argv[2]
00121     
00122     model_img = cv.LoadImage(model_name)
00123     cand_img = cv.LoadImage(candidate)
00124     
00125     for i in range(5):
00126         thres = .8 - (i * .1)
00127         print 'thres %.2f' % thres
00128         joint_viz = match_images(model_img, cand_img, thres)
00129         win_name = 'surf%.2f' % thres
00130         cv.NamedWindow(win_name, 0)
00131         cv.ShowImage(win_name, joint_viz)
00132     
00133     while not rospy.is_shutdown():
00134         cv.WaitKey(10)
00135 
00136 ###############
00137 ##############
00138 
00139 if __name__ == '__main__':
00140     mode = 'image'
00141     #if mode = 'image':
00142     #    find pose of model
00143     #    find normal of model
00144     #    record angles of features.
00145 
00146     if mode=='image':
00147         test_thresholds()
00148 
00149     if mode=='live':
00150         model_name = sys.argv[1]
00151         model_img = cv.LoadImage(model_name)
00152         model_gray = fea.grayscale(model_img)
00153         msurf_loc, msurf_desc = fea.surf(model_gray)
00154         prosilica = rc.Prosilica('prosilica', 'streaming')
00155         cv.NamedWindow('surf', 1)
00156         si = ShowImage('surf')
00157         si.start()
00158 
00159         #Each feature is a row in db
00160         features_db = sp.KDTree(np.array(msurf_desc))
00161         #pdb.set_trace()
00162  
00163         while not rospy.is_shutdown():
00164             print '..'
00165             image = prosilica.get_frame()
00166             print 'saving image'
00167             cv.SaveImage('frame.png', image)
00168             print '>'
00169             img_gray = fea.grayscale(image)
00170             locs, descs = fea.surf(img_gray)
00171             match_idxs = []
00172             for i, desc in enumerate(descs):
00173                 dists, idxs = features_db.query(np.array(desc), 2)
00174                 ratio = dists[0] / dists[1]
00175                 if ratio < .49:
00176                     match_idxs.append(i)
00177             img_viz = fea.draw_surf(image, locs, (255,0,0))
00178             img_viz = fea.draw_surf(img_viz, [locs[i] for i in match_idxs], (0,0,255))
00179             si.image = img_viz
00180             print '%d matches found' % len(match_idxs)
00181             #print len(desc), desc.__class__, len(descs[0]), descs[0].__class__
00182             #si.image = image
00183 
00184    
00185 
00186 
00187 
00188 
00189 
00190 
00191 
00192 
00193 
00194 
00195 
00196 
00197 
00198 
00199 


hai_sandbox
Author(s): Hai Nguyen
autogenerated on Wed Nov 27 2013 11:46:56