| 12345678910111213141516171819202122232425262728293031323334 |
- from PIL import Image
- import torch
- import cv2
- from romatch import roma_outdoor
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- if torch.backends.mps.is_available():
- device = torch.device('mps')
- if __name__ == "__main__":
- from argparse import ArgumentParser
- parser = ArgumentParser()
- parser.add_argument("--im_A_path", default="assets/sacre_coeur_A.jpg", type=str)
- parser.add_argument("--im_B_path", default="assets/sacre_coeur_B.jpg", type=str)
- args, _ = parser.parse_known_args()
- im1_path = args.im_A_path
- im2_path = args.im_B_path
- # Create model
- roma_model = roma_outdoor(device=device)
- W_A, H_A = Image.open(im1_path).size
- W_B, H_B = Image.open(im2_path).size
- # Match
- warp, certainty = roma_model.match(im1_path, im2_path, device=device)
- # Sample matches for estimation
- matches, certainty = roma_model.sample(warp, certainty)
- kpts1, kpts2 = roma_model.to_pixel_coordinates(matches, H_A, W_A, H_B, W_B)
- F, mask = cv2.findFundamentalMat(
- kpts1.cpu().numpy(), kpts2.cpu().numpy(), ransacReprojThreshold=0.2, method=cv2.USAC_MAGSAC, confidence=0.999999, maxIters=10000
- )
|