added flann matching to find_obj.py

samples/python2/find_obj.py

@@ -1,14 +1,19 @@
 import numpy as np
 import cv2
 from common import anorm
+from functools import partial
 
 help_message = '''SURF image match 
 
 USAGE: findobj.py [ <image1> <image2> ]
 '''
 
+FLANN_INDEX_KDTREE = 1  # bug: flann enums are missing
 
-def match(desc1, desc2, r_threshold = 0.75):
+flann_params = dict(algorithm = FLANN_INDEX_KDTREE,
+                    trees = 4)
+
+def match_bruteforce(desc1, desc2, r_threshold = 0.75):
     res = []
     for i in xrange(len(desc1)):
         dist = anorm( desc2 - desc1[i] )
@@ -18,6 +23,14 @@ def match(desc1, desc2, r_threshold = 0.75):
             res.append((i, n1))
     return np.array(res)
 
+def match_flann(desc1, desc2, r_threshold = 0.6):
+    flann = cv2.flann_Index(desc2, flann_params)
+    idx2, dist = flann.knnSearch(desc1, 2, params = {}) # bug: need to provide empty dict
+    mask = dist[:,0] / dist[:,1] < r_threshold
+    idx1 = np.arange(len(desc1))
+    pairs = np.int32( zip(idx1, idx2[:,0]) )
+    return pairs[mask]
+
 def draw_match(img1, img2, p1, p2, status = None, H = None):
     h1, w1 = img1.shape[:2]
     h2, w2 = img2.shape[:2]
@@ -50,6 +63,7 @@ def draw_match(img1, img2, p1, p2, status = None, H = None):
             cv2.line(vis, (x2+w1-r, y2+r), (x2+w1+r, y2-r), col, thickness)
     return vis
 
+
 if __name__ == '__main__':
     import sys
     try: fn1, fn2 = sys.argv[1:3]
@@ -68,12 +82,21 @@ if __name__ == '__main__':
     desc2.shape = (-1, surf.descriptorSize())
     print 'img1 - %d features, img2 - %d features' % (len(kp1), len(kp2))
 
-    m = match(desc1, desc2)
-    matched_p1 = np.array([kp1[i].pt for i, j in m])
-    matched_p2 = np.array([kp2[j].pt for i, j in m])
-    H, status = cv2.findHomography(matched_p1, matched_p2, cv2.RANSAC, 10.0)
-    print '%d / %d  inliers/matched' % (np.sum(status), len(status))
+    def match_and_draw(match, r_threshold):
+        m = match(desc1, desc2, r_threshold)
+        matched_p1 = np.array([kp1[i].pt for i, j in m])
+        matched_p2 = np.array([kp2[j].pt for i, j in m])
+        H, status = cv2.findHomography(matched_p1, matched_p2, cv2.RANSAC, 5.0)
+        print '%d / %d  inliers/matched' % (np.sum(status), len(status))
+
+        vis = draw_match(img1, img2, matched_p1, matched_p2, status, H)
+        return vis
 
-    vis = draw_match(img1, img2, matched_p1, matched_p2, status, H)
-    cv2.imshow('find_obj SURF', vis)
+    print 'bruteforce match:',
+    vis_brute = match_and_draw( match_bruteforce, 0.75 )
+    print 'flann match:',
+    vis_flann = match_and_draw( match_flann, 0.6 ) # flann tends to find more distant second
+                                                   # neighbours, so r_threshold is decreased
+    cv2.imshow('find_obj SURF', vis_brute)
+    cv2.imshow('find_obj SURF flann', vis_flann)
     cv2.waitKey()