Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,338 Bytes
4c954ae |
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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
import argparse
import os
from os.path import join
import sys
import cv2
import torch
from matplotlib import pyplot as plt
from tqdm import tqdm
from gluestick import batch_to_np, numpy_image_to_torch, GLUESTICK_ROOT
from gluestick.drawing import plot_images, plot_lines, plot_color_line_matches, plot_keypoints, plot_matches
# from gluestick.models.two_view_pipeline import TwoViewPipeline
from line_matching.two_view_pipeline import TwoViewPipeline
from scalelsd.base import show, WireframeGraph
def main():
# Parse input parameters
parser = argparse.ArgumentParser(
prog='GlueStick Demo',
description='Demo app to show the point and line matches obtained by GlueStick')
parser.add_argument('-inum', default=None, type=int)
parser.add_argument('-imax', default=None, type=int)
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('--model', default='scalelsd', type=str)
parser.add_argument('--test_root', type=str, default='data-ssl/0images-pre/')
args = parser.parse_args()
# Evaluation config
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,
# 'merge_line_endpoints': False,
},
'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()
pipeline_model.extractor.update_conf(None)
md = args.model
root = args.test_root
if args.inum is not None:
ids = [args.inum]
elif args.imax is not None:
ids = range(args.inum, args.imax+1)
else:
l_imgs = int(len(os.listdir(root))/2)
ids = range(l_imgs)
for id in tqdm(ids):
saveto = f'temp_output/matching_results/{md}/{id}'
os.makedirs(saveto, exist_ok=True)
args.img1 = root + f'ref_{str(id)}.png'
args.img2 = root + f'tgt_{str(id)}.png'
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]
# Plot the matches
gray0 = cv2.imread(args.img1, 0)
gray1 = cv2.imread(args.img2, 0)
img0, img1 = cv2.cvtColor(gray0, cv2.COLOR_GRAY2BGR), cv2.cvtColor(gray1, cv2.COLOR_GRAY2BGR)
plot_images([img0, img1], 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.tight_layout()
plt.savefig(f'{saveto}/{md}_det_{id}.png')
plot_images([img0, img1], dpi=200, pad=2.0)
plot_color_line_matches([matched_lines0, matched_lines1], lw=3)
plt.gcf().canvas.manager.set_window_title('Line Matches')
# plt.tight_layout()
plt.savefig(f'{saveto}/{md}_mat_{id}.png')
whitebg = 1
show.Canvas.white_overlay = whitebg
painter = show.painters.HAWPainter()
fig_file = f'{saveto}/{md}_det1.png'
outputs = {'lines_pred': line_seg0.reshape(-1,4)}
with show.image_canvas(args.img1, fig_file=fig_file) as ax:
# painter.draw_wireframe(ax,outputs, edge_color='orange', vertex_color='Cyan')
painter.draw_wireframe(ax,outputs, edge_color='midnightblue', vertex_color='deeppink')
fig_file = f'{saveto}/{md}_det2.png'
outputs = {'lines_pred': line_seg1.reshape(-1,4)}
with show.image_canvas(args.img2, fig_file=fig_file) as ax:
# painter.draw_wireframe(ax,outputs, edge_color='orange', vertex_color='Cyan')
painter.draw_wireframe(ax,outputs, edge_color='midnightblue', vertex_color='deeppink')
if __name__ == '__main__':
main()
|