feature_detector.py 1.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. from collections import OrderedDict
  2. import cv2 as cv
  3. class FeatureDetector:
  4. DETECTOR_CHOICES = OrderedDict()
  5. try:
  6. cv.xfeatures2d_SURF.create() # check if the function can be called
  7. DETECTOR_CHOICES['surf'] = cv.xfeatures2d_SURF.create
  8. except (AttributeError, cv.error):
  9. print("SURF not available")
  10. # if SURF not available, ORB is default
  11. DETECTOR_CHOICES['orb'] = cv.ORB.create
  12. try:
  13. DETECTOR_CHOICES['sift'] = cv.SIFT_create
  14. except AttributeError:
  15. print("SIFT not available")
  16. try:
  17. DETECTOR_CHOICES['brisk'] = cv.BRISK_create
  18. except AttributeError:
  19. print("BRISK not available")
  20. try:
  21. DETECTOR_CHOICES['akaze'] = cv.AKAZE_create
  22. except AttributeError:
  23. print("AKAZE not available")
  24. DEFAULT_DETECTOR = list(DETECTOR_CHOICES.keys())[0]
  25. def __init__(self, detector=DEFAULT_DETECTOR, **kwargs):
  26. self.detector = FeatureDetector.DETECTOR_CHOICES[detector](**kwargs)
  27. def detect_features(self, img, *args, **kwargs):
  28. return cv.detail.computeImageFeatures2(self.detector, img,
  29. *args, **kwargs)
  30. @staticmethod
  31. def draw_keypoints(img, features, **kwargs):
  32. kwargs.setdefault('color', (0, 255, 0))
  33. keypoints = features.getKeypoints()
  34. return cv.drawKeypoints(img, keypoints, None, **kwargs)