|
from pathlib import Path |
|
import argparse |
|
import cv2 |
|
import matplotlib.cm as cm |
|
import torch |
|
import numpy as np |
|
from utils.nnmatching import NNMatching |
|
from utils.misc import (AverageTimer, VideoStreamer, make_matching_plot_fast, frame2tensor) |
|
|
|
torch.set_grad_enabled(False) |
|
|
|
|
|
def compute_essential(matched_kp1, matched_kp2, K): |
|
pts1 = cv2.undistortPoints(matched_kp1,cameraMatrix=K, distCoeffs = (-0.117918271740560,0.075246403574314,0,0)) |
|
pts2 = cv2.undistortPoints(matched_kp2,cameraMatrix=K, distCoeffs = (-0.117918271740560,0.075246403574314,0,0)) |
|
K_1 = np.eye(3) |
|
|
|
ransac_model, ransac_inliers = cv2.findEssentialMat(pts1, pts2, K_1, method=cv2.RANSAC, prob=0.999, threshold=0.001, maxIters=10000) |
|
if ransac_inliers is None or ransac_model.shape != (3,3): |
|
ransac_inliers = np.array([]) |
|
ransac_model = None |
|
return ransac_model, ransac_inliers, pts1, pts2 |
|
|
|
|
|
sizer = (960, 640) |
|
focallength_x = 4.504986436499113e+03/(6744/sizer[0]) |
|
focallength_y = 4.513311442889859e+03/(4502/sizer[1]) |
|
K = np.eye(3) |
|
K[0,0] = focallength_x |
|
K[1,1] = focallength_y |
|
K[0,2] = 3.363322177533149e+03/(6744/sizer[0]) |
|
K[1,2] = 2.291824660547715e+03/(4502/sizer[1]) |
|
|
|
|
|
if __name__ == '__main__': |
|
parser = argparse.ArgumentParser( |
|
description='DarkFeat demo', |
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter) |
|
parser.add_argument( |
|
'--input', type=str, |
|
help='path to an image directory') |
|
parser.add_argument( |
|
'--output_dir', type=str, default=None, |
|
help='Directory where to write output frames (If None, no output)') |
|
|
|
parser.add_argument( |
|
'--image_glob', type=str, nargs='+', default=['*.ARW'], |
|
help='Glob if a directory of images is specified') |
|
parser.add_argument( |
|
'--resize', type=int, nargs='+', default=[640, 480], |
|
help='Resize the input image before running inference. If two numbers, ' |
|
'resize to the exact dimensions, if one number, resize the max ' |
|
'dimension, if -1, do not resize') |
|
parser.add_argument( |
|
'--force_cpu', action='store_true', |
|
help='Force pytorch to run in CPU mode.') |
|
parser.add_argument('--model_path', type=str, |
|
help='Path to the pretrained model') |
|
|
|
opt = parser.parse_args() |
|
print(opt) |
|
|
|
assert len(opt.resize) == 2 |
|
print('Will resize to {}x{} (WxH)'.format(opt.resize[0], opt.resize[1])) |
|
|
|
device = 'cuda' if torch.cuda.is_available() and not opt.force_cpu else 'cpu' |
|
print('Running inference on device \"{}\"'.format(device)) |
|
matching = NNMatching(opt.model_path).eval().to(device) |
|
keys = ['keypoints', 'scores', 'descriptors'] |
|
|
|
vs = VideoStreamer(opt.input, opt.resize, opt.image_glob) |
|
frame, ret = vs.next_frame() |
|
assert ret, 'Error when reading the first frame (try different --input?)' |
|
|
|
frame_tensor = frame2tensor(frame, device) |
|
last_data = matching.darkfeat({'image': frame_tensor}) |
|
last_data = {k+'0': [last_data[k]] for k in keys} |
|
last_data['image0'] = frame_tensor |
|
last_frame = frame |
|
last_image_id = 0 |
|
|
|
if opt.output_dir is not None: |
|
print('==> Will write outputs to {}'.format(opt.output_dir)) |
|
Path(opt.output_dir).mkdir(exist_ok=True) |
|
|
|
timer = AverageTimer() |
|
|
|
while True: |
|
frame, ret = vs.next_frame() |
|
if not ret: |
|
print('Finished demo_darkfeat.py') |
|
break |
|
timer.update('data') |
|
stem0, stem1 = last_image_id, vs.i - 1 |
|
|
|
frame_tensor = frame2tensor(frame, device) |
|
pred = matching({**last_data, 'image1': frame_tensor}) |
|
kpts0 = last_data['keypoints0'][0].cpu().numpy() |
|
kpts1 = pred['keypoints1'][0].cpu().numpy() |
|
matches = pred['matches0'][0].cpu().numpy() |
|
confidence = pred['matching_scores0'][0].cpu().numpy() |
|
timer.update('forward') |
|
|
|
valid = matches > -1 |
|
mkpts0 = kpts0[valid] |
|
mkpts1 = kpts1[matches[valid]] |
|
|
|
E, inliers, pts1, pts2 = compute_essential(mkpts0, mkpts1, K) |
|
color = cm.jet(np.clip(confidence[valid][inliers[:, 0].astype('bool')] * 2 - 1, -1, 1)) |
|
|
|
text = [ |
|
'DarkFeat', |
|
'Matches: {}'.format(inliers.sum()) |
|
] |
|
|
|
out = make_matching_plot_fast( |
|
last_frame, frame, mkpts0[inliers[:, 0].astype('bool')], mkpts1[inliers[:, 0].astype('bool')], color, text, |
|
path=None, small_text=' ') |
|
|
|
if opt.output_dir is not None: |
|
stem = 'matches_{:06}_{:06}'.format(stem0, stem1) |
|
out_file = str(Path(opt.output_dir, stem + '.png')) |
|
print('Writing image to {}'.format(out_file)) |
|
cv2.imwrite(out_file, out) |
|
|
