Spaces:
Running
Running
File size: 1,946 Bytes
6ba5875 c74a070 6ba5875 c74a070 6ba5875 c74a070 6ba5875 c74a070 6ba5875 c74a070 6ba5875 c74a070 6ba5875 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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
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"
)
|