Spaces:
Running
Running
File size: 2,356 Bytes
13760e8 |
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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import os
import sys
import torch
import numpy as np
import math
import cv2
os.environ['CUDA_VISIBLE_DEVICES']='1'
from models.liftfeat_wrapper import LiftFeat,MODEL_PATH
import argparse
parser=argparse.ArgumentParser(description='HPatch dataset evaluation script')
parser.add_argument('--name',type=str,default='LiftFeat',help='experiment name')
parser.add_argument('--img1',type=str,default='./assert/ref.jpg',help='reference image path')
parser.add_argument('--img2',type=str,default='./assert/query.jpg',help='query image path')
parser.add_argument('--gpu',type=str,default='0',help='GPU ID')
args=parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
def warp_corners_and_draw_matches(ref_points, dst_points, img1, img2):
# Calculate the Homography matrix
H, mask = cv2.findHomography(ref_points, dst_points, cv2.USAC_MAGSAC, 3.5, maxIters=1_000, confidence=0.999)
mask = mask.flatten()
# Get corners of the first image (image1)
h, w = img1.shape[:2]
corners_img1 = np.array([[0, 0], [w-1, 0], [w-1, h-1], [0, h-1]], dtype=np.float32).reshape(-1, 1, 2)
# Warp corners to the second image (image2) space
warped_corners = cv2.perspectiveTransform(corners_img1, H)
# Draw the warped corners in image2
img2_with_corners = img2.copy()
# Prepare keypoints and matches for drawMatches function
keypoints1 = [cv2.KeyPoint(float(p[0]), float(p[1]), 5) for p in ref_points]
keypoints2 = [cv2.KeyPoint(float(p[0]), float(p[1]), 5) for p in dst_points]
matches = [cv2.DMatch(i,i,0) for i in range(len(mask)) if mask[i]]
# Draw inlier matches
img_matches = cv2.drawMatches(img1, keypoints1, img2_with_corners, keypoints2, matches, None,
matchColor=(0, 255, 0), flags=2)
return img_matches
if __name__=="__main__":
liftfeat=LiftFeat(weight=MODEL_PATH,detect_threshold=0.05)
img1=cv2.imread(args.img1)
img2=cv2.imread(args.img2)
# import pdb;pdb.set_trace()
mkpts1,mkpts2=liftfeat.match_liftfeat(img1,img2)
canvas=warp_corners_and_draw_matches(mkpts1,mkpts2,img1,img2)
import matplotlib.pyplot as plt
plt.figure(figsize=[12,12])
plt.imshow(canvas[...,::-1])
plt.savefig(os.path.join(os.path.dirname(__file__),'match.jpg'), dpi=300, bbox_inches='tight')
plt.show()
|