|
import sys,os |
|
|
|
from skimage import io |
|
import urllib.request |
|
import numpy as np |
|
import cv2 |
|
#from matplotlib import pyplot as plt |
|
|
|
dataPath = "C:\\to\\your\datadir" |
|
|
|
class AutoBanDectector(): |
|
MIN_MATCH_COUNT = 30 |
|
def __init__(self): |
|
# init sift and surf |
|
self.sift = cv2.xfeatures2d.SIFT_create() |
|
self.surf = cv2.xfeatures2d.SURF_create() |
|
# build sift result dict |
|
# todo: pickle it |
|
self.resDict = {} |
|
|
|
for fn in os.listdir(dataPath+"\\times"): |
|
pfn = dataPath+"\\times\\"+fn |
|
nt = fn.split(".")[0] |
|
img = cv2.imread(pfn,0) |
|
self.resDict[nt] = self.sift.detectAndCompute(img,None) |
|
# init flann searcher |
|
FLANN_INDEX_KDTREE = 0 |
|
index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5) |
|
search_params = dict(checks = 50) |
|
self.flann = cv2.FlannBasedMatcher(index_params, search_params) |
|
# build surf points for template |
|
img = cv2.imread(dataPath+"\\tem2x.png",0) # todo: from structed data |
|
self.kpTemplate, self.desTemplate = self.surf.detectAndCompute(img,None) |
|
def fuck(self,fn): |
|
img2d = io.imread(fn,0) |
|
if None is img2d: |
|
raise Exception("bad file %s"%fn) |
|
sca = max(img2d.shape)//1024 |
|
while sca > 1: |
|
sca = sca>>1 |
|
img2d = cv2.pyrDown(img2d) |
|
|
|
kp, des = self.surf.detectAndCompute(img2d,None) |
|
|
|
matches = self.flann.knnMatch(self.desTemplate,des,k=2) |
|
|
|
# from ocv offical document |
|
good = [] |
|
for m,n in matches: |
|
if m.distance < 0.7*n.distance: |
|
good.append(m) |
|
print("found",len(good)) |
|
if len(good)>self.MIN_MATCH_COUNT: |
|
# warp perspective |
|
src_pts = np.float32([ self.kpTemplate[m.queryIdx].pt for m in good ]).reshape(-1,1,2) |
|
dst_pts = np.float32([ kp[m.trainIdx].pt for m in good ]).reshape(-1,1,2) |
|
imgText = cv2.warpPerspective(img2d, cv2.findHomography(dst_pts, src_pts, cv2.RANSAC)[0], (img2d.shape[0],img2d.shape[1])) |
|
imgText = imgText[0:75, 65:101] # todo: from structed data |
|
return self.findBestmatch(imgText) |
|
else: |
|
# fail |
|
return None |
|
def findBestmatch(self,img): |
|
_,des = self.sift.detectAndCompute(img,None) |
|
maxMatch = "none" |
|
maxMatchNum = 0 |
|
for k in self.resDict.keys(): |
|
matches = self.flann.knnMatch(des,self.resDict[k][1],k=2) |
|
mt = 0 |
|
for m,n in matches: |
|
if m.distance < 0.7*n.distance: |
|
mt+=1 |
|
#print(k,mt) |
|
if mt>=maxMatchNum: |
|
maxMatch = k |
|
maxMatchNum = mt |
|
#print("bestmatch",maxMatch,maxMatchNum) |
|
return maxMatch |
|
|
|
if __name__ == "__main__": |
|
detector = AutoBanDectector() |
|
print(detector.fuck(sys.argv[1])) |