47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
import cv2
|
|
import os
|
|
from matplotlib import pyplot as plt
|
|
|
|
def FLANN():
|
|
|
|
targetPath = 'AI.jpg'
|
|
trainingImage = cv2.imread(targetPath, flags=0)
|
|
|
|
templatePath = 'image/'
|
|
icons = os.listdir(templatePath)
|
|
iconMatch= dict({'name': '未识别', 'value': 0})
|
|
for icon in icons:
|
|
queryImage = cv2.imread(templatePath + icon, 0)
|
|
|
|
sift = cv2.SIFT_create()
|
|
kp1, des1 = sift.detectAndCompute(queryImage, None)
|
|
kp2, des2 = sift.detectAndCompute(trainingImage, None)
|
|
|
|
indexParams = dict(algorithm=0, trees=5)
|
|
searchParams = dict(checks=50)
|
|
flann = cv2.FlannBasedMatcher(indexParams,searchParams)
|
|
matches = flann.knnMatch(des1,des2,k=2)
|
|
|
|
matchesMask = [[0,0] for i in range(len(matches))]
|
|
matchNumber = 0
|
|
|
|
for i,(m,n) in enumerate(matches):
|
|
if m.distance < 0.7 * n.distance:
|
|
matchesMask[i] = [1,0]
|
|
matchNumber = matchNumber+1
|
|
|
|
drawParams = dict(matchColor = (0,255,0), matchesMask = matchesMask[:50], flags = 0)
|
|
resultImage = cv2.drawMatchesKnn(queryImage,kp1,trainingImage,kp2,matches[:50],None,**drawParams)
|
|
|
|
if matchNumber > iconMatch['value']:
|
|
iconMatch['name'] = icon.split('_')[0]
|
|
iconMatch['value'] = matchNumber
|
|
|
|
return resultImage, iconMatch
|
|
|
|
if __name__ == '__main__':
|
|
resultImage, res = FLANN()
|
|
# plt.imshow(resultImage)
|
|
# plt.show()
|
|
print(resultImage, res)
|