demo_fundamental.py 1.2 KB

12345678910111213141516171819202122232425262728293031323334
  1. from PIL import Image
  2. import torch
  3. import cv2
  4. from romatch import roma_outdoor
  5. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  6. if torch.backends.mps.is_available():
  7. device = torch.device('mps')
  8. if __name__ == "__main__":
  9. from argparse import ArgumentParser
  10. parser = ArgumentParser()
  11. parser.add_argument("--im_A_path", default="assets/sacre_coeur_A.jpg", type=str)
  12. parser.add_argument("--im_B_path", default="assets/sacre_coeur_B.jpg", type=str)
  13. args, _ = parser.parse_known_args()
  14. im1_path = args.im_A_path
  15. im2_path = args.im_B_path
  16. # Create model
  17. roma_model = roma_outdoor(device=device)
  18. W_A, H_A = Image.open(im1_path).size
  19. W_B, H_B = Image.open(im2_path).size
  20. # Match
  21. warp, certainty = roma_model.match(im1_path, im2_path, device=device)
  22. # Sample matches for estimation
  23. matches, certainty = roma_model.sample(warp, certainty)
  24. kpts1, kpts2 = roma_model.to_pixel_coordinates(matches, H_A, W_A, H_B, W_B)
  25. F, mask = cv2.findFundamentalMat(
  26. kpts1.cpu().numpy(), kpts2.cpu().numpy(), ransacReprojThreshold=0.2, method=cv2.USAC_MAGSAC, confidence=0.999999, maxIters=10000
  27. )