You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

47 lines
1.4KB

  1. import cv2
  2. import os
  3. from matplotlib import pyplot as plt
  4. def FLANN():
  5. targetPath = 'AI.jpg'
  6. trainingImage = cv2.imread(targetPath, flags=0)
  7. templatePath = 'image/'
  8. icons = os.listdir(templatePath)
  9. iconMatch= dict({'name': '未识别', 'value': 0})
  10. for icon in icons:
  11. queryImage = cv2.imread(templatePath + icon, 0)
  12. sift = cv2.SIFT_create()
  13. kp1, des1 = sift.detectAndCompute(queryImage, None)
  14. kp2, des2 = sift.detectAndCompute(trainingImage, None)
  15. indexParams = dict(algorithm=0, trees=5)
  16. searchParams = dict(checks=50)
  17. flann = cv2.FlannBasedMatcher(indexParams,searchParams)
  18. matches = flann.knnMatch(des1,des2,k=2)
  19. matchesMask = [[0,0] for i in range(len(matches))]
  20. matchNumber = 0
  21. for i,(m,n) in enumerate(matches):
  22. if m.distance < 0.7 * n.distance:
  23. matchesMask[i] = [1,0]
  24. matchNumber = matchNumber+1
  25. drawParams = dict(matchColor = (0,255,0), matchesMask = matchesMask[:50], flags = 0)
  26. resultImage = cv2.drawMatchesKnn(queryImage,kp1,trainingImage,kp2,matches[:50],None,**drawParams)
  27. if matchNumber > iconMatch['value']:
  28. iconMatch['name'] = icon.split('_')[0]
  29. iconMatch['value'] = matchNumber
  30. return resultImage, iconMatch
  31. if __name__ == '__main__':
  32. resultImage, res = FLANN()
  33. # plt.imshow(resultImage)
  34. # plt.show()
  35. print(resultImage, res)