import torch from DeDoDe import dedode_detector_L, dedode_descriptor_B from DeDoDe.matchers.dual_softmax_matcher import DualSoftMaxMatcher from DeDoDe.utils import * from PIL import Image import cv2 def draw_matches(im_A, kpts_A, im_B, kpts_B): kpts_A = [cv2.KeyPoint(x, y, 1.0) for x, y in kpts_A.cpu().numpy()] kpts_B = [cv2.KeyPoint(x, y, 1.0) for x, y in kpts_B.cpu().numpy()] matches_A_to_B = [cv2.DMatch(idx, idx, 0.0) for idx in range(len(kpts_A))] im_A, im_B = np.array(im_A), np.array(im_B) ret = cv2.drawMatches(im_A, kpts_A, im_B, kpts_B, matches_A_to_B, None) return ret detector = dedode_detector_L(weights=torch.load("dedode_detector_L.pth")) descriptor = dedode_descriptor_B(weights=torch.load("dedode_descriptor_B.pth")) matcher = DualSoftMaxMatcher() im_A_path = "assets/im_A.jpg" im_B_path = "assets/im_B.jpg" im_A = Image.open(im_A_path) im_B = Image.open(im_B_path) W_A, H_A = im_A.size W_B, H_B = im_B.size detections_A = detector.detect_from_path(im_A_path, num_keypoints=10_000) keypoints_A, P_A = detections_A["keypoints"], detections_A["confidence"] detections_B = detector.detect_from_path(im_B_path, num_keypoints=10_000) keypoints_B, P_B = detections_B["keypoints"], detections_B["confidence"] description_A = descriptor.describe_keypoints_from_path(im_A_path, keypoints_A)[ "descriptions" ] description_B = descriptor.describe_keypoints_from_path(im_B_path, keypoints_B)[ "descriptions" ] matches_A, matches_B, batch_ids = matcher.match( keypoints_A, description_A, keypoints_B, description_B, P_A=P_A, P_B=P_B, normalize=True, inv_temp=20, threshold=0.1, ) # Increasing threshold -> fewer matches, fewer outliers matches_A, matches_B = matcher.to_pixel_coords(matches_A, matches_B, H_A, W_A, H_B, W_B) import cv2 import numpy as np Image.fromarray(draw_matches(im_A, matches_A[::5], im_B, matches_B[::5])).save( "demo/matches.png" )