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
00035
00036
00037
00038
00039
00040
00041
00042
00043
00044
00045
00046
00047
00048
00049
00050
00051
00052
00053
00054
00055
00056
00057
00058
00059
00060
00061
00062
00063
00064
00065
00066
00067
00068
00069
00070
00071
00072
00073 def match_images(model_img, cand_img, threshold=.8):
00074
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
00090 if ratio < threshold:
00091 matched.append((i, idxs[0]))
00092
00093
00094
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
00107
00108
00109
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
00142
00143
00144
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
00160 features_db = sp.KDTree(np.array(msurf_desc))
00161
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
00182
00183
00184
00185
00186
00187
00188
00189
00190
00191
00192
00193
00194
00195
00196
00197
00198
00199