feature_matcher.py 3.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. import math
  2. import cv2 as cv
  3. import numpy as np
  4. class FeatureMatcher:
  5. MATCHER_CHOICES = ('homography', 'affine')
  6. DEFAULT_MATCHER = 'homography'
  7. DEFAULT_RANGE_WIDTH = -1
  8. def __init__(self,
  9. matcher_type=DEFAULT_MATCHER,
  10. range_width=DEFAULT_RANGE_WIDTH,
  11. **kwargs):
  12. if matcher_type == "affine":
  13. """https://docs.opencv.org/4.x/d3/dda/classcv_1_1detail_1_1AffineBestOf2NearestMatcher.html""" # noqa
  14. self.matcher = cv.detail_AffineBestOf2NearestMatcher(**kwargs)
  15. elif range_width == -1:
  16. """https://docs.opencv.org/4.x/d4/d26/classcv_1_1detail_1_1BestOf2NearestMatcher.html""" # noqa
  17. self.matcher = cv.detail_BestOf2NearestMatcher(**kwargs)
  18. else:
  19. """https://docs.opencv.org/4.x/d8/d72/classcv_1_1detail_1_1BestOf2NearestRangeMatcher.html""" # noqa
  20. self.matcher = cv.detail_BestOf2NearestRangeMatcher(
  21. range_width, **kwargs
  22. )
  23. def match_features(self, features, *args, **kwargs):
  24. pairwise_matches = self.matcher.apply2(features, *args, **kwargs)
  25. self.matcher.collectGarbage()
  26. return pairwise_matches
  27. @staticmethod
  28. def draw_matches_matrix(imgs, features, matches, conf_thresh=1,
  29. inliers=False, **kwargs):
  30. matches_matrix = FeatureMatcher.get_matches_matrix(matches)
  31. for idx1, idx2 in FeatureMatcher.get_all_img_combinations(len(imgs)):
  32. match = matches_matrix[idx1, idx2]
  33. if match.confidence < conf_thresh:
  34. continue
  35. if inliers:
  36. kwargs['matchesMask'] = match.getInliers()
  37. yield idx1, idx2, FeatureMatcher.draw_matches(
  38. imgs[idx1], features[idx1],
  39. imgs[idx2], features[idx2],
  40. match,
  41. **kwargs
  42. )
  43. @staticmethod
  44. def draw_matches(img1, features1, img2, features2, match1to2, **kwargs):
  45. kwargs.setdefault('flags', cv.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)
  46. keypoints1 = features1.getKeypoints()
  47. keypoints2 = features2.getKeypoints()
  48. matches = match1to2.getMatches()
  49. return cv.drawMatches(
  50. img1, keypoints1, img2, keypoints2, matches, None, **kwargs
  51. )
  52. @staticmethod
  53. def get_matches_matrix(pairwise_matches):
  54. return FeatureMatcher.array_in_sqare_matrix(pairwise_matches)
  55. @staticmethod
  56. def get_confidence_matrix(pairwise_matches):
  57. matches_matrix = FeatureMatcher.get_matches_matrix(pairwise_matches)
  58. match_confs = [[m.confidence for m in row] for row in matches_matrix]
  59. match_conf_matrix = np.array(match_confs)
  60. return match_conf_matrix
  61. @staticmethod
  62. def array_in_sqare_matrix(array):
  63. matrix_dimension = int(math.sqrt(len(array)))
  64. rows = []
  65. for i in range(0, len(array), matrix_dimension):
  66. rows.append(array[i:i+matrix_dimension])
  67. return np.array(rows)
  68. def get_all_img_combinations(number_imgs):
  69. ii, jj = np.triu_indices(number_imgs, k=1)
  70. for i, j in zip(ii, jj):
  71. yield i, j
  72. @staticmethod
  73. def get_match_conf(match_conf, feature_detector_type):
  74. if match_conf is None:
  75. match_conf = \
  76. FeatureMatcher.get_default_match_conf(feature_detector_type)
  77. return match_conf
  78. @staticmethod
  79. def get_default_match_conf(feature_detector_type):
  80. if feature_detector_type == 'orb':
  81. return 0.3
  82. return 0.65