Jiaming Sun
4 years ago
committed by
Jiaming Sun
parent
63818b0287
commit
8a6eabeaa3
4 changed files with 334 additions and 14 deletions
@ -0,0 +1,240 @@ |
||||
front_matter = """ |
||||
------------------------------------------------------------------------ |
||||
Online demo for [LoFTR](https://zju3dv.github.io/loftr/). |
||||
|
||||
This demo is heavily inspired by [SuperGlue](https://github.com/magicleap/SuperGluePretrainedNetwork/). |
||||
We thank the authors for their execellent work. |
||||
------------------------------------------------------------------------ |
||||
""" |
||||
|
||||
import os |
||||
import argparse |
||||
from pathlib import Path |
||||
import cv2 |
||||
import torch |
||||
import numpy as np |
||||
import matplotlib.cm as cm |
||||
|
||||
os.sys.path.append("../") # Add the project directory |
||||
from src.loftr import LoFTR, default_cfg |
||||
from src.config.default import get_cfg_defaults |
||||
try: |
||||
from demo.utils import (AverageTimer, VideoStreamer, |
||||
make_matching_plot_fast, make_matching_plot, frame2tensor) |
||||
except: |
||||
raise ImportError("This demo requires utils.py from SuperGlue, please use run_demo.sh to start this script.") |
||||
|
||||
|
||||
torch.set_grad_enabled(False) |
||||
|
||||
if __name__ == '__main__': |
||||
parser = argparse.ArgumentParser( |
||||
description='LoFTR online demo', |
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter) |
||||
parser.add_argument('--weight', type=str, help="Path to the checkpoint.") |
||||
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( |
||||
'--no_display', action='store_true', |
||||
help='Do not display images to screen. Useful if running remotely') |
||||
parser.add_argument( |
||||
'--save_video', action='store_true', |
||||
help='Save output (with match visualizations) to a video.') |
||||
parser.add_argument( |
||||
'--save_input', action='store_true', |
||||
help='Save the input images to a video (for gathering repeatable input source).') |
||||
parser.add_argument( |
||||
'--skip_frames', type=int, default=1, |
||||
help="Skip frames from webcam input.") |
||||
parser.add_argument( |
||||
'--top_k', type=int, default=2000, help="The max vis_range (please refer to the code).") |
||||
parser.add_argument( |
||||
'--bottom_k', type=int, default=0, help="The min vis_range (please refer to the code).") |
||||
|
||||
opt = parser.parse_args() |
||||
print(front_matter) |
||||
parser.print_help() |
||||
|
||||
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') |
||||
|
||||
if torch.cuda.is_available(): |
||||
device = 'cuda' |
||||
else: |
||||
raise RuntimeError("GPU is required to run this demo.") |
||||
|
||||
# Initialize LoFTR |
||||
matcher = LoFTR(config=default_cfg) |
||||
matcher.load_state_dict(torch.load(opt.weight)['state_dict']) |
||||
matcher = matcher.eval().to(device=device) |
||||
|
||||
# Configure I/O |
||||
if opt.save_video: |
||||
print('Writing video to loftr-matches.mp4...') |
||||
writer = cv2.VideoWriter('loftr-matches.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 15, (640*2 + 10, 480)) |
||||
if opt.save_input: |
||||
print('Writing video to demo-input.mp4...') |
||||
input_writer = cv2.VideoWriter('demo-input.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 15, (640, 480)) |
||||
|
||||
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_id = 0 |
||||
last_image_id = 0 |
||||
frame_tensor = frame2tensor(frame, device) |
||||
last_data = {'image0': frame_tensor} |
||||
last_frame = frame |
||||
|
||||
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: |
||||
window_name = 'LoFTR Matches' |
||||
cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) |
||||
cv2.resizeWindow(window_name, (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 reference image (left)\n' |
||||
'\td/f: move the range of the matches (ranked by confidence) to visualize\n' |
||||
'\tc/v: increase/decrease the length of the visualization range (i.e., total number of matches) to show\n' |
||||
'\tq: quit') |
||||
|
||||
timer = AverageTimer() |
||||
vis_range = [opt.bottom_k, opt.top_k] |
||||
|
||||
while True: |
||||
frame_id += 1 |
||||
frame, ret = vs.next_frame() |
||||
if frame_id % opt.skip_frames != 0: |
||||
# print("Skipping frame.") |
||||
continue |
||||
if opt.save_input: |
||||
inp = np.stack([frame]*3, -1) |
||||
inp_rgb = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB) |
||||
input_writer.write(inp_rgb) |
||||
if not ret: |
||||
print('Finished demo_loftr.py') |
||||
break |
||||
timer.update('data') |
||||
stem0, stem1 = last_image_id, vs.i - 1 |
||||
|
||||
frame_tensor = frame2tensor(frame, device) |
||||
last_data = {**last_data, 'image1': frame_tensor} |
||||
matcher(last_data) |
||||
|
||||
total_n_matches = len(last_data['mkpts0_f']) |
||||
mkpts0 = last_data['mkpts0_f'].cpu().numpy()[vis_range[0]:vis_range[1]] |
||||
mkpts1 = last_data['mkpts1_f'].cpu().numpy()[vis_range[0]:vis_range[1]] |
||||
mconf = last_data['mconf'].cpu().numpy()[vis_range[0]:vis_range[1]] |
||||
|
||||
# Normalize confidence. |
||||
if len(mconf) > 0: |
||||
conf_vis_min = 0. |
||||
conf_min = mconf.min() |
||||
conf_max = mconf.max() |
||||
mconf = (mconf - conf_vis_min) / (conf_max - conf_vis_min + 1e-5) |
||||
|
||||
timer.update('forward') |
||||
alpha = 0 |
||||
color = cm.jet(mconf, alpha=alpha) |
||||
|
||||
text = [ |
||||
f'LoFTR', |
||||
'# Matches (showing/total): {}/{}'.format(len(mkpts0), total_n_matches), |
||||
] |
||||
small_text = [ |
||||
f'Showing matches from {vis_range[0]}:{vis_range[1]}', |
||||
f'Confidence Range: {conf_min:.2f}:{conf_max:.2f}', |
||||
'Image Pair: {:06}:{:06}'.format(stem0, stem1), |
||||
] |
||||
out = make_matching_plot_fast( |
||||
last_frame, frame, mkpts0, mkpts1, mkpts0, mkpts1, color, text, |
||||
path=None, show_keypoints=False, small_text=small_text) |
||||
|
||||
# Save high quality png, optionally with dynamic alpha support (unreleased yet). |
||||
# save_path = 'demo_vid/{:06}'.format(frame_id) |
||||
# make_matching_plot( |
||||
# last_frame, frame, mkpts0, mkpts1, mkpts0, mkpts1, color, text, |
||||
# path=save_path, show_keypoints=opt.show_keypoints, small_text=small_text) |
||||
|
||||
if not opt.no_display: |
||||
if opt.save_video: |
||||
writer.write(out) |
||||
cv2.imshow('LoFTR Matches', out) |
||||
key = chr(cv2.waitKey(1) & 0xFF) |
||||
if key == 'q': |
||||
if opt.save_video: |
||||
writer.release() |
||||
if opt.save_input: |
||||
input_writer.release() |
||||
vs.cleanup() |
||||
print('Exiting...') |
||||
break |
||||
elif key == 'n': |
||||
last_data['image0'] = frame_tensor |
||||
last_frame = frame |
||||
last_image_id = (vs.i - 1) |
||||
frame_id_left = frame_id |
||||
elif key in ['d', 'f']: |
||||
if key == 'd': |
||||
if vis_range[0] >= 0: |
||||
vis_range[0] -= 200 |
||||
vis_range[1] -= 200 |
||||
if key =='f': |
||||
vis_range[0] += 200 |
||||
vis_range[1] += 200 |
||||
print(f'\nChanged the vis_range to {vis_range[0]}:{vis_range[1]}') |
||||
elif key in ['c', 'v']: |
||||
if key == 'c': |
||||
vis_range[1] -= 50 |
||||
if key =='v': |
||||
vis_range[1] += 50 |
||||
print(f'\nChanged the vis_range[1] to {vis_range[1]}') |
||||
elif opt.output_dir is not None: |
||||
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) |
||||
else: |
||||
raise ValueError("output_dir is required when no display is given.") |
||||
timer.update('viz') |
||||
timer.print() |
||||
|
||||
|
||||
cv2.destroyAllWindows() |
||||
vs.cleanup() |
||||
@ -0,0 +1,34 @@ |
||||
#!/bin/bash |
||||
set -e |
||||
# set -x |
||||
|
||||
if [ ! -f utils.py ]; then |
||||
echo "Downloading utils.py from the SuperGlue repo." |
||||
echo "We cannot provide this file directly due to its strict licence." |
||||
wget https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/master/models/utils.py |
||||
fi |
||||
|
||||
# Use webcam 0 as input source. |
||||
input=0 |
||||
# or use a pre-recorded video given the path. |
||||
# input=/home/sunjiaming/Downloads/scannet_test/$scene_name.mp4 |
||||
|
||||
# Toggle indoor/outdoor model here. |
||||
model_ckpt=../weights/indoor_ds.ckpt |
||||
# model_ckpt=../weights/outdoor_ds.ckpt |
||||
|
||||
# Optionally assign the GPU ID. |
||||
# export CUDA_VISIBLE_DEVICES=0 |
||||
|
||||
echo "Running LoFTR demo.." |
||||
eval "$(conda shell.zsh hook)" |
||||
conda activate loftr |
||||
python demo_loftr.py --weight $model_ckpt --input $input |
||||
# To save the input video and output match visualizations. |
||||
# python demo_loftr.py --weight $model_ckpt --input $input --save_video --save_input |
||||
|
||||
# Running on remote GPU servers with no GUI. |
||||
# Save images first. |
||||
# python demo_loftr.py --weight $model_ckpt --input $input --no_display --output_dir="./demo_images/" |
||||
# Then convert them to a video. |
||||
# ffmpeg -framerate 15 -pattern_type glob -i '*.png' -c:v libx264 -r 30 -pix_fmt yuv420p out.mp4 |
||||
Loading…
Reference in new issue