opencv.py

import numpy as np
import cv2
import os
import fnmatch

FLANN_INDEX_KDTREE = 1  # bug: flann enums are missing

flann_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 10)


def match_flann(desc1, desc2, r_threshold):
    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]
    vis = np.zeros((max(h1, h2), w1+w2), np.uint8)
    vis[:h1, :w1] = img1
    vis[:h2, w1:w1+w2] = img2
    vis = cv2.cvtColor(vis, cv2.COLOR_GRAY2BGR)

    if status is None:
        status = np.ones(len(p1), np.bool_)
    green = (0, 128, 0)
    red = (0, 0, 128)
    for (x1, y1), (x2, y2), inlier in zip(np.int32(p1), np.int32(p2), status):
        col = [red, green][inlier]
        if inlier:
            cv2.line(vis, (x1, y1), (x2+w1, y2), col)
            cv2.circle(vis, (x1, y1), 2, col, -1)
            cv2.circle(vis, (x2+w1, y2), 2, col, -1)
        else:
            r = 2
            thickness = 3
            cv2.line(vis, (x1-r, y1-r), (x1+r, y1+r), col, thickness)
            cv2.line(vis, (x1-r, y1+r), (x1+r, y1-r), col, thickness)
            cv2.line(vis, (x2+w1-r, y2-r), (x2+w1+r, y2+r), col, thickness)
            cv2.line(vis, (x2+w1-r, y2+r), (x2+w1+r, y2-r), col, thickness)

    if H is not None:
        corners = np.float32([[0, 0], [w1, 0], [w1, h1], [0, h1]])
        corners = np.int32( cv2.perspectiveTransform(corners.reshape(1, -1, 2), H).reshape(-1, 2) + (w1, 0) )
        cv2.polylines(vis, [corners], True, (255, 0, 0), 2)

    return vis


def test(fn1, fn2) :
    img1 = cv2.imread(fn1, 0)
    img2 = cv2.imread(fn2, 0)

    surf = cv2.SURF(400)
    kp1, desc1 = surf.detectAndCompute(img1, None)
    kp2, desc2 = surf.detectAndCompute(img2, None)
    desc1.shape = (-1, surf.descriptorSize())
    desc2.shape = (-1, surf.descriptorSize())
    #print 'img1 - %d features, img2 - %d features' % (len(kp1), len(kp2))

    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])

        print 'img1 - %d matches, img2 - %d matches' % (len(matched_p1), len(matched_p2))

        if matched_p1.size == 0 :
            return None
        if matched_p2.size == 0 :
            return None

        rows, cols = matched_p1.shape
        if (rows < 4 and cols < 4) :
            return None

        rows, cols = matched_p2.shape
        if (rows < 4 and cols < 4) :
            return None

        H, status = cv2.findHomography(matched_p1, matched_p2, cv2.RANSAC, 6)
        #print "\n%d / %d  inliers/matched" % (np.sum(status), len(status))

        vis = draw_match(img1, img2, matched_p1, matched_p2, status, H)
        return vis

    #print 'flann match:',
    vis_flann = match_and_draw( match_flann, 0.8)
    cv2.imshow('find_obj SURF flann', vis_flann)
    if vis_flann is not None :
#        #shutil.copy2(fn2, '/Users/joshua/Desktop/found')
        0xFF & cv2.waitKey()



path = "/Users/joshua/Desktop/found"

matches = []
for root, dirnames, filenames in os.walk(path):
    for filename in fnmatch.filter(filenames, '*'):
        matches.append(os.path.join(root, filename))

fn1 = "/Users/joshua/Desktop/tag.png"
for match in matches :
    print match, "\n"
    try:
        test(fn1, match)
    except:
        pass