added flann matching to find_obj.py

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