camera_estimator.py 906 B

123456789101112131415161718192021222324252627
  1. from collections import OrderedDict
  2. import cv2 as cv
  3. import numpy as np
  4. from .stitching_error import StitchingError
  5. class CameraEstimator:
  6. CAMERA_ESTIMATOR_CHOICES = OrderedDict()
  7. CAMERA_ESTIMATOR_CHOICES['homography'] = cv.detail_HomographyBasedEstimator
  8. CAMERA_ESTIMATOR_CHOICES['affine'] = cv.detail_AffineBasedEstimator
  9. DEFAULT_CAMERA_ESTIMATOR = list(CAMERA_ESTIMATOR_CHOICES.keys())[0]
  10. def __init__(self, estimator=DEFAULT_CAMERA_ESTIMATOR, **kwargs):
  11. self.estimator = CameraEstimator.CAMERA_ESTIMATOR_CHOICES[estimator](
  12. **kwargs
  13. )
  14. def estimate(self, features, pairwise_matches):
  15. b, cameras = self.estimator.apply(features, pairwise_matches, None)
  16. if not b:
  17. raise StitchingError("Homography estimation failed.")
  18. for cam in cameras:
  19. cam.R = cam.R.astype(np.float32)
  20. return cameras