#! /usr/bin/env python3 # # %BANNER_BEGIN% # --------------------------------------------------------------------- # %COPYRIGHT_BEGIN% # # Magic Leap, Inc. ("COMPANY") CONFIDENTIAL # # Unpublished Copyright (c) 2020 # Magic Leap, Inc., All Rights Reserved. # # NOTICE: All information contained herein is, and remains the property # of COMPANY. The intellectual and technical concepts contained herein # are proprietary to COMPANY and may be covered by U.S. and Foreign # Patents, patents in process, and are protected by trade secret or # copyright law. Dissemination of this information or reproduction of # this material is strictly forbidden unless prior written permission is # obtained from COMPANY. Access to the source code contained herein is # hereby forbidden to anyone except current COMPANY employees, managers # or contractors who have executed Confidentiality and Non-disclosure # agreements explicitly covering such access. # # The copyright notice above does not evidence any actual or intended # publication or disclosure of this source code, which includes # information that is confidential and/or proprietary, and is a trade # secret, of COMPANY. ANY REPRODUCTION, MODIFICATION, DISTRIBUTION, # PUBLIC PERFORMANCE, OR PUBLIC DISPLAY OF OR THROUGH USE OF THIS # SOURCE CODE WITHOUT THE EXPRESS WRITTEN CONSENT OF COMPANY IS # STRICTLY PROHIBITED, AND IN VIOLATION OF APPLICABLE LAWS AND # INTERNATIONAL TREATIES. THE RECEIPT OR POSSESSION OF THIS SOURCE # CODE AND/OR RELATED INFORMATION DOES NOT CONVEY OR IMPLY ANY RIGHTS # TO REPRODUCE, DISCLOSE OR DISTRIBUTE ITS CONTENTS, OR TO MANUFACTURE, # USE, OR SELL ANYTHING THAT IT MAY DESCRIBE, IN WHOLE OR IN PART. # # %COPYRIGHT_END% # ---------------------------------------------------------------------- # %AUTHORS_BEGIN% # # Originating Authors: Paul-Edouard Sarlin # Daniel DeTone # Tomasz Malisiewicz # # %AUTHORS_END% # --------------------------------------------------------------------*/ # %BANNER_END% from pathlib import Path import argparse import cv2 import matplotlib.cm as cm import torch from models.matching import Matching from models.utils import (AverageTimer, VideoStreamer, make_matching_plot_fast, frame2tensor) torch.set_grad_enabled(False) if __name__ == '__main__': parser = argparse.ArgumentParser( description='SuperGlue demo', formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument( '--input', type=str, default='0', help='ID of a USB webcam, URL of an IP camera, ' 'or path to an image directory or movie file') 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=['*.png', '*.jpg', '*.jpeg'], help='Glob if a directory of images is specified') parser.add_argument( '--skip', type=int, default=1, help='Images to skip if input is a movie or directory') parser.add_argument( '--max_length', type=int, default=1000000, help='Maximum length if input is a movie or directory') 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( '--superglue', choices={'indoor', 'outdoor'}, default='indoor', help='SuperGlue weights') parser.add_argument( '--max_keypoints', type=int, default=-1, help='Maximum number of keypoints detected by Superpoint' ' (\'-1\' keeps all keypoints)') parser.add_argument( '--keypoint_threshold', type=float, default=0.005, help='SuperPoint keypoint detector confidence threshold') parser.add_argument( '--nms_radius', type=int, default=4, help='SuperPoint Non Maximum Suppression (NMS) radius' ' (Must be positive)') parser.add_argument( '--sinkhorn_iterations', type=int, default=20, help='Number of Sinkhorn iterations performed by SuperGlue') parser.add_argument( '--match_threshold', type=float, default=0.2, help='SuperGlue match threshold') parser.add_argument( '--show_keypoints', action='store_true', help='Show the detected keypoints') parser.add_argument( '--no_display', action='store_true', help='Do not display images to screen. Useful if running remotely') parser.add_argument( '--force_cpu', action='store_true', help='Force pytorch to run in CPU mode.') opt = parser.parse_args() print(opt) if len(opt.resize) == 2 and opt.resize[1] == -1: opt.resize = opt.resize[0:1] if len(opt.resize) == 2: print('Will resize to {}x{} (WxH)'.format( opt.resize[0], opt.resize[1])) elif len(opt.resize) == 1 and opt.resize[0] > 0: print('Will resize max dimension to {}'.format(opt.resize[0])) elif len(opt.resize) == 1: print('Will not resize images') else: raise ValueError('Cannot specify more than two integers for --resize') device = 'cuda' if torch.cuda.is_available() and not opt.force_cpu else 'cpu' print('Running inference on device \"{}\"'.format(device)) config = { 'superpoint': { 'nms_radius': opt.nms_radius, 'keypoint_threshold': opt.keypoint_threshold, 'max_keypoints': opt.max_keypoints }, 'superglue': { 'weights': opt.superglue, 'sinkhorn_iterations': opt.sinkhorn_iterations, 'match_threshold': opt.match_threshold, } } matching = Matching(config).eval().to(device) keys = ['keypoints', 'scores', 'descriptors'] vs = VideoStreamer(opt.input, opt.resize, opt.skip, opt.image_glob, opt.max_length) frame, ret = vs.next_frame() assert ret, 'Error when reading the first frame (try different --input?)' frame_tensor = frame2tensor(frame, device) last_data = matching.superpoint({'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) # Create a window to display the demo. if not opt.no_display: cv2.namedWindow('SuperGlue matches', cv2.WINDOW_NORMAL) cv2.resizeWindow('SuperGlue matches', 640*2, 480) else: print('Skipping visualization, will not show a GUI.') # Print the keyboard help menu. print('==> Keyboard control:\n' '\tn: select the current frame as the anchor\n' '\te/r: increase/decrease the keypoint confidence threshold\n' '\td/f: increase/decrease the match filtering threshold\n' '\tk: toggle the visualization of keypoints\n' '\tq: quit') timer = AverageTimer() while True: frame, ret = vs.next_frame() if not ret: print('Finished demo_superglue.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]] color = cm.jet(confidence[valid]) text = [ 'SuperGlue', 'Keypoints: {}:{}'.format(len(kpts0), len(kpts1)), 'Matches: {}'.format(len(mkpts0)) ] k_thresh = matching.superpoint.config['keypoint_threshold'] m_thresh = matching.superglue.config['match_threshold'] small_text = [ 'Keypoint Threshold: {:.4f}'.format(k_thresh), 'Match Threshold: {:.2f}'.format(m_thresh), 'Image Pair: {:06}:{:06}'.format(stem0, stem1), ] out = make_matching_plot_fast( last_frame, frame, kpts0, kpts1, mkpts0, mkpts1, color, text, path=None, show_keypoints=opt.show_keypoints, small_text=small_text) if not opt.no_display: cv2.imshow('SuperGlue matches', out) key = chr(cv2.waitKey(1) & 0xFF) if key == 'q': vs.cleanup() print('Exiting (via q) demo_superglue.py') break elif key == 'n': # set the current frame as anchor last_data = {k+'0': pred[k+'1'] for k in keys} last_data['image0'] = frame_tensor last_frame = frame last_image_id = (vs.i - 1) elif key in ['e', 'r']: # Increase/decrease keypoint threshold by 10% each keypress. d = 0.1 * (-1 if key == 'e' else 1) matching.superpoint.config['keypoint_threshold'] = min(max( 0.0001, matching.superpoint.config['keypoint_threshold']*(1+d)), 1) print('\nChanged the keypoint threshold to {:.4f}'.format( matching.superpoint.config['keypoint_threshold'])) elif key in ['d', 'f']: # Increase/decrease match threshold by 0.05 each keypress. d = 0.05 * (-1 if key == 'd' else 1) matching.superglue.config['match_threshold'] = min(max( 0.05, matching.superglue.config['match_threshold']+d), .95) print('\nChanged the match threshold to {:.2f}'.format( matching.superglue.config['match_threshold'])) elif key == 'k': opt.show_keypoints = not opt.show_keypoints timer.update('viz') timer.print() if opt.output_dir is not None: #stem = 'matches_{:06}_{:06}'.format(last_image_id, vs.i-1) stem = 'matches_{:06}_{:06}'.format(stem0, stem1) out_file = str(Path(opt.output_dir, stem + '.png')) print('\nWriting image to {}'.format(out_file)) cv2.imwrite(out_file, out) cv2.destroyAllWindows() vs.cleanup()