00001 import roslib; roslib.load_manifest('hai_sandbox')
00002 import cv
00003 import numpy as np
00004 import scipy.spatial as sp
00005 import math
00006
00007 class SURFMatcher:
00008 def __init__(self):
00009 self.model_images = {}
00010 self.model_fea = {}
00011
00012 def add_file(self, model_name, label):
00013 model_img = cv.LoadImage(model_name)
00014 self.add_model(model_img, label)
00015
00016 def add_model(self, model_img, label):
00017 mgray = grayscale(model_img)
00018 m_loc, m_desc = surf(mgray)
00019 self.model_images[label] = model_img
00020 self.model_fea[label] = {'loc': m_loc, 'desc': m_desc}
00021
00022 def build_db(self):
00023 fea_l = []
00024 labels_l = []
00025 locs_l = []
00026 for k in self.model_fea:
00027 fea_l.append(np.array(self.model_fea[k]['desc']))
00028 locs_l.append(np.array(self.model_fea[k]['loc']))
00029 labels_l.append(np.array([k for i in range(len(self.model_fea[k]['desc']))]))
00030
00031 self.labels = np.row_stack(labels_l)
00032 self.locs = np.row_stack(locs_l)
00033 self.tree = sp.KDTree(np.row_stack(fea_l))
00034
00035 def match(self, desc, thres=.6):
00036 dists, idxs = self.tree.query(np.array(desc), 2)
00037 ratio = dists[0] / dists[1]
00038 if ratio < threshold:
00039 desc = self.tree.data[idxs[0]]
00040 loc = self.locs[idxs[0]]
00041 return desc, loc
00042 else:
00043 return None
00044
00045 def concat_images(a, b):
00046 img_height = max(a.height, b.height)
00047 c = cv.CreateImage((a.width+b.width, img_height), a.depth, a.channels)
00048 a_area = cv.GetSubRect(c, (0,0, a.width, a.height))
00049 b_area = cv.GetSubRect(c, (a.width, 0, b.width, b.height))
00050 cv.Add(a, a_area, a_area)
00051 cv.Add(b, b_area, b_area)
00052 return c
00053
00054 def clone(something):
00055 if something.__class__ == cv.cvmat:
00056 return cv.CloneMat(something)
00057 else:
00058 return cv.CloneImage(something)
00059
00060 def draw_surf(image, keypoints, color):
00061 rimage = clone(image)
00062
00063 for loc, lap, size, d, hess in keypoints:
00064 loc = tuple(np.array(np.round(loc), dtype='int').tolist())
00065 circ_rad = int(round(size/4.))
00066 cv.Circle(rimage, loc, circ_rad, color, 1, cv.CV_AA)
00067 cv.Circle(rimage, loc, 2, color, -1, cv.CV_AA)
00068
00069 drad = math.radians(d)
00070 line_len = circ_rad
00071 loc_end = (np.matrix(np.round( circ_rad * np.matrix([np.cos(drad), np.sin(drad)]).T + np.matrix(loc).T), dtype='int')).A1.tolist()
00072 cv.Line(rimage, loc, tuple(loc_end), color, thickness=1, lineType=cv.CV_AA)
00073
00074 return rimage
00075
00076 def draw_surf2(image, keypoints, colors):
00077 rimage = clone(image)
00078 for i, k in enumerate(keypoints):
00079 loc, lap, size, d, hess = k
00080 loc = tuple(np.array(np.round(loc), dtype='int').tolist())
00081 c = tuple(np.matrix(colors[:,i],dtype='int').T.A1)
00082 color = (int(c[0]), int(c[1]), int(c[2]))
00083
00084 cv.Circle(rimage, loc, 5, color, 1, cv.CV_AA)
00085 return rimage
00086
00087 def draw_harris(image, keypoints, color):
00088 rimage = clone(image)
00089 for loc in keypoints:
00090 loc = tuple(np.array(np.round(loc), dtype='int').tolist())
00091 cv.Circle(rimage, loc, 5, color, 1, cv.CV_AA)
00092 return rimage
00093
00094 def draw_star(image, keypoints, color):
00095 rimage = clone(image)
00096 color_arr = np.array(color)
00097 max_resp = - 999999
00098 min_resp = 999999
00099 for _, _, response in keypoints:
00100 max_resp = max(response, max_resp)
00101 min_resp = min(response, min_resp)
00102 range_resp = max_resp - min_resp
00103
00104 for loc, size, response in keypoints:
00105 loc = tuple(np.array(np.round(loc), dtype='int').tolist())
00106 color_weight = ((response - min_resp) / range_resp)
00107 c = tuple((color_weight * color_arr).tolist())
00108 cv.Circle(rimage, loc, int(round(size/2.0)), c, 1, cv.CV_AA)
00109 return rimage
00110
00111
00112 def star(image):
00113 star_stor = cv.CreateMemStorage()
00114 star_keypoints = cv.GetStarKeypoints(image, star_stor)
00115 del star_stor
00116 return star_keypoints
00117
00118
00119
00120
00121 def surf(image_gray, params=(1, 3000,3,4)):
00122 surf_stor = cv.CreateMemStorage()
00123 surf_r = cv.ExtractSURF(image_gray, None, surf_stor, params)
00124 del surf_stor
00125 return surf_r
00126
00127
00128
00129 def surf_color(image, params=(1,3000,3,4)):
00130 gray = grayscale(image)
00131 return surf(gray, params)
00132
00133
00134
00135 def harris(image_gray):
00136 eig_image = cv.CreateImage(cv.GetSize(image_gray), cv.IPL_DEPTH_32F, 1)
00137 temp_image = cv.CreateImage(cv.GetSize(image_gray), cv.IPL_DEPTH_32F, 1)
00138 return cv.GoodFeaturesToTrack(image_gray, eig_image, temp_image, 300, .1, 1.0, useHarris = True)
00139
00140 def grayscale(image):
00141 image_gray = cv.CreateImage(cv.GetSize(image), cv.IPL_DEPTH_8U,1)
00142 cv.CvtColor(image, image_gray, cv.CV_BGR2GRAY)
00143 return image_gray