Spaces:
Running
Running
File size: 1,151 Bytes
c608946 c74a070 c608946 c74a070 c608946 c74a070 c608946 c74a070 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
from PIL import Image
import torch
import cv2
from roma import roma_outdoor
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
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,
)
|