Spaces:
Running
Running
| 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(): | |
| # 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("-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() | |
| # 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, | |
| }, | |
| "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] | |
| # Plot the matches | |
| 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() | |