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)