|
import argparse |
|
import os |
|
from os.path import join |
|
|
|
import cv2 |
|
import torch |
|
from matplotlib import pyplot as plt |
|
|
|
from gluestick import batch_to_np, numpy_image_to_torch, GLUESTICK_ROOT |
|
from .drawing import plot_images, plot_lines, plot_color_line_matches, plot_keypoints, plot_matches |
|
from .models.two_view_pipeline import TwoViewPipeline |
|
|
|
|
|
def main(): |
|
|
|
parser = argparse.ArgumentParser( |
|
prog='GlueStick Demo', |
|
description='Demo app to show the point and line matches obtained by GlueStick') |
|
parser.add_argument('-img1', default=join('resources' + os.path.sep + 'img1.jpg')) |
|
parser.add_argument('-img2', default=join('resources' + os.path.sep + 'img2.jpg')) |
|
parser.add_argument('--max_pts', type=int, default=1000) |
|
parser.add_argument('--max_lines', type=int, default=300) |
|
parser.add_argument('--skip-imshow', default=False, action='store_true') |
|
args = parser.parse_args() |
|
|
|
|
|
conf = { |
|
'name': 'two_view_pipeline', |
|
'use_lines': True, |
|
'extractor': { |
|
'name': 'wireframe', |
|
'sp_params': { |
|
'force_num_keypoints': False, |
|
'max_num_keypoints': args.max_pts, |
|
}, |
|
'wireframe_params': { |
|
'merge_points': True, |
|
'merge_line_endpoints': True, |
|
}, |
|
'max_n_lines': args.max_lines, |
|
}, |
|
'matcher': { |
|
'name': 'gluestick', |
|
'weights': str(GLUESTICK_ROOT / 'resources' / 'weights' / 'checkpoint_GlueStick_MD.tar'), |
|
'trainable': False, |
|
}, |
|
'ground_truth': { |
|
'from_pose_depth': False, |
|
} |
|
} |
|
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
pipeline_model = TwoViewPipeline(conf).to(device).eval() |
|
|
|
gray0 = cv2.imread(args.img1, 0) |
|
gray1 = cv2.imread(args.img2, 0) |
|
|
|
torch_gray0, torch_gray1 = numpy_image_to_torch(gray0), numpy_image_to_torch(gray1) |
|
torch_gray0, torch_gray1 = torch_gray0.to(device)[None], torch_gray1.to(device)[None] |
|
x = {'image0': torch_gray0, 'image1': torch_gray1} |
|
pred = pipeline_model(x) |
|
|
|
pred = batch_to_np(pred) |
|
kp0, kp1 = pred["keypoints0"], pred["keypoints1"] |
|
m0 = pred["matches0"] |
|
|
|
line_seg0, line_seg1 = pred["lines0"], pred["lines1"] |
|
line_matches = pred["line_matches0"] |
|
|
|
valid_matches = m0 != -1 |
|
match_indices = m0[valid_matches] |
|
matched_kps0 = kp0[valid_matches] |
|
matched_kps1 = kp1[match_indices] |
|
|
|
valid_matches = line_matches != -1 |
|
match_indices = line_matches[valid_matches] |
|
matched_lines0 = line_seg0[valid_matches] |
|
matched_lines1 = line_seg1[match_indices] |
|
|
|
|
|
img0, img1 = cv2.cvtColor(gray0, cv2.COLOR_GRAY2BGR), cv2.cvtColor(gray1, cv2.COLOR_GRAY2BGR) |
|
plot_images([img0, img1], ['Image 1 - detected lines', 'Image 2 - detected lines'], dpi=200, pad=2.0) |
|
plot_lines([line_seg0, line_seg1], ps=4, lw=2) |
|
plt.gcf().canvas.manager.set_window_title('Detected Lines') |
|
plt.savefig('detected_lines.png') |
|
|
|
plot_images([img0, img1], ['Image 1 - detected points', 'Image 2 - detected points'], dpi=200, pad=2.0) |
|
plot_keypoints([kp0, kp1], colors='c') |
|
plt.gcf().canvas.manager.set_window_title('Detected Points') |
|
plt.savefig('detected_points.png') |
|
|
|
plot_images([img0, img1], ['Image 1 - line matches', 'Image 2 - line matches'], dpi=200, pad=2.0) |
|
plot_color_line_matches([matched_lines0, matched_lines1], lw=2) |
|
plt.gcf().canvas.manager.set_window_title('Line Matches') |
|
plt.savefig('line_matches.png') |
|
|
|
plot_images([img0, img1], ['Image 1 - point matches', 'Image 2 - point matches'], dpi=200, pad=2.0) |
|
plot_matches(matched_kps0, matched_kps1, 'green', lw=1, ps=0) |
|
plt.gcf().canvas.manager.set_window_title('Point Matches') |
|
plt.savefig('detected_points.png') |
|
if not args.skip_imshow: |
|
plt.show() |
|
|
|
|
|
if __name__ == '__main__': |
|
main() |
|
|
