gim-online / third_party /RoRD /extractMatch.py
Vincentqyw
update: rord
ab830e5
raw
history blame
5.5 kB
import argparse
import numpy as np
import imageio
import torch
from tqdm import tqdm
import time
import scipy
import scipy.io
import scipy.misc
import os
import sys
from lib.model_test import D2Net
from lib.utils import preprocess_image
from lib.pyramid import process_multiscale
import cv2
import matplotlib.pyplot as plt
from PIL import Image
from skimage.feature import match_descriptors
from skimage.measure import ransac
from skimage.transform import ProjectiveTransform, AffineTransform
import pydegensac
parser = argparse.ArgumentParser(description='Feature extraction script')
parser.add_argument('imgs', type=str, nargs=2)
parser.add_argument(
'--preprocessing', type=str, default='caffe',
help='image preprocessing (caffe or torch)'
)
parser.add_argument(
'--model_file', type=str,
help='path to the full model'
)
parser.add_argument(
'--no-relu', dest='use_relu', action='store_false',
help='remove ReLU after the dense feature extraction module'
)
parser.set_defaults(use_relu=True)
parser.add_argument(
'--sift', dest='use_sift', action='store_true',
help='Show sift matching as well'
)
parser.set_defaults(use_sift=False)
def extract(image, args, model, device):
if len(image.shape) == 2:
image = image[:, :, np.newaxis]
image = np.repeat(image, 3, -1)
input_image = preprocess_image(
image,
preprocessing=args.preprocessing
)
with torch.no_grad():
keypoints, scores, descriptors = process_multiscale(
torch.tensor(
input_image[np.newaxis, :, :, :].astype(np.float32),
device=device
),
model,
scales=[1]
)
keypoints = keypoints[:, [1, 0, 2]]
feat = {}
feat['keypoints'] = keypoints
feat['scores'] = scores
feat['descriptors'] = descriptors
return feat
def rordMatching(image1, image2, feat1, feat2, matcher="BF"):
if(matcher == "BF"):
t0 = time.time()
bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True)
matches = bf.match(feat1['descriptors'], feat2['descriptors'])
matches = sorted(matches, key=lambda x:x.distance)
t1 = time.time()
print("Time to extract matches: ", t1-t0)
print("Number of raw matches:", len(matches))
match1 = [m.queryIdx for m in matches]
match2 = [m.trainIdx for m in matches]
keypoints_left = feat1['keypoints'][match1, : 2]
keypoints_right = feat2['keypoints'][match2, : 2]
np.random.seed(0)
t0 = time.time()
H, inliers = pydegensac.findHomography(keypoints_left, keypoints_right, 10.0, 0.99, 10000)
t1 = time.time()
print("Time for ransac: ", t1-t0)
n_inliers = np.sum(inliers)
print('Number of inliers: %d.' % n_inliers)
inlier_keypoints_left = [cv2.KeyPoint(point[0], point[1], 1) for point in keypoints_left[inliers]]
inlier_keypoints_right = [cv2.KeyPoint(point[0], point[1], 1) for point in keypoints_right[inliers]]
placeholder_matches = [cv2.DMatch(idx, idx, 1) for idx in range(n_inliers)]
draw_params = dict(matchColor = (0,255,0),
singlePointColor = (255,0,0),
# matchesMask = matchesMask,
flags = 0)
image3 = cv2.drawMatches(image1, inlier_keypoints_left, image2, inlier_keypoints_right, placeholder_matches, None, **draw_params)
plt.figure(figsize=(20, 20))
plt.imshow(image3)
plt.axis('off')
plt.show()
def siftMatching(img1, img2):
img1 = np.array(cv2.cvtColor(np.array(img1), cv2.COLOR_BGR2RGB))
img2 = np.array(cv2.cvtColor(np.array(img2), cv2.COLOR_BGR2RGB))
# surf = cv2.xfeatures2d.SURF_create(100)
surf = cv2.xfeatures2d.SIFT_create()
kp1, des1 = surf.detectAndCompute(img1, None)
kp2, des2 = surf.detectAndCompute(img2, None)
FLANN_INDEX_KDTREE = 0
index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5)
search_params = dict(checks = 50)
flann = cv2.FlannBasedMatcher(index_params, search_params)
matches = flann.knnMatch(des1,des2,k=2)
good = []
for m, n in matches:
if m.distance < 0.7*n.distance:
good.append(m)
src_pts = np.float32([ kp1[m.queryIdx].pt for m in good ]).reshape(-1, 2)
dst_pts = np.float32([ kp2[m.trainIdx].pt for m in good ]).reshape(-1, 2)
model, inliers = pydegensac.findHomography(src_pts, dst_pts, 10.0, 0.99, 10000)
n_inliers = np.sum(inliers)
print('Number of inliers: %d.' % n_inliers)
inlier_keypoints_left = [cv2.KeyPoint(point[0], point[1], 1) for point in src_pts[inliers]]
inlier_keypoints_right = [cv2.KeyPoint(point[0], point[1], 1) for point in dst_pts[inliers]]
placeholder_matches = [cv2.DMatch(idx, idx, 1) for idx in range(n_inliers)]
image3 = cv2.drawMatches(img1, inlier_keypoints_left, img2, inlier_keypoints_right, placeholder_matches, None)
cv2.imshow('Matches', image3)
cv2.waitKey(0)
src_pts = np.float32([ inlier_keypoints_left[m.queryIdx].pt for m in placeholder_matches ]).reshape(-1, 2)
dst_pts = np.float32([ inlier_keypoints_right[m.trainIdx].pt for m in placeholder_matches ]).reshape(-1, 2)
return src_pts, dst_pts
if __name__ == '__main__':
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
args = parser.parse_args()
model = D2Net(
model_file=args.model_file,
use_relu=args.use_relu,
use_cuda=use_cuda
)
image1 = np.array(Image.open(args.imgs[0]))
image2 = np.array(Image.open(args.imgs[1]))
print('--\nRoRD\n--')
feat1 = extract(image1, args, model, device)
feat2 = extract(image2, args, model, device)
print("Features extracted.")
rordMatching(image1, image2, feat1, feat2, matcher="BF")
if(args.use_sift):
print('--\nSIFT\n--')
siftMatching(image1, image2)