赞
踩
本脚本是一个基于Python的应用,旨在演示如何使用SuperGlue算法进行图像之间的特征匹配。SuperGlue是一个强大的特征匹配工具,能够在不同的图像之间找到对应的关键点。这个工具尤其适用于计算机视觉任务,如立体视觉、图像拼接、对象识别和追踪等场景。脚本使用PyTorch框架,并且可以选择在CPU或GPU上运行。
脚本的工作流程如下:
- 解析命令行参数,用于设置输入输出目录、图像尺寸、SuperGlue配置等。
- 根据用户选择,决定算法是在CPU还是GPU上执行。
- 加载预设的配置,初始化SuperPoint和SuperGlue模型。
- 定义图像预处理函数来调整图像大小。
- 加载两幅图像,调整它们的大小,并将它们转换为PyTorch张量。
- 使用SuperPoint提取关键点和描述符。
- 使用SuperGlue算法匹配两幅图像的关键点。
- 可视化并打印匹配的关键点坐标。
- 如果设置了输出目录,将结果图像写到磁盘上。
这个脚本展示了如何在实践中使用深度学习模型来处理实际问题,并提供了图像匹配演示。
- #! /usr/bin/env python3
- import argparse
- import matplotlib.cm as cm
- import cv2
- from pathlib import Path
- import torch
- from models.matching import Matching
- from models.utils import (make_matching_plot_fast, frame2tensor)
- torch.set_grad_enabled(False) # 关闭PyTorch的梯度计算,提高效率,因为我们不需要进行模型训练
-
- # 创建命令行参数解析器,以便从命令行接收参数
- parser = argparse.ArgumentParser(description='SuperGlue',
- formatter_class=argparse.ArgumentDefaultsHelpFormatter)
- # 添加命令行参数
- parser.add_argument('--input', type=str, default='assets/freiburg_sequence/',
- help='Input directory or video file')
- parser.add_argument('--output_dir', type=str, default=None,
- help='Directory to write output frames (default: None)')
- parser.add_argument('--resize', type=int, nargs='+', default=[1241, 376],
- help='Resize input frames (default: [640, 480])')
- parser.add_argument('--superglue', choices={'indoor', 'outdoor'}, default='outdoor',
- help='SuperGlue weights (default: indoor)')
- parser.add_argument('--show_keypoints', action='store_true',
- help='Show detected keypoints (default: False)')
- parser.add_argument('--no_display', action='store_true',
- help='Do not display images (useful when running remotely)')
- parser.add_argument('--force_cpu', action='store_true',
- help='Force PyTorch to run on CPU')
-
- # 解析命令行参数
- opt = parser.parse_args()
-
- # 确定程序是运行在GPU还是CPU
- device = 'cuda' if torch.cuda.is_available() and not opt.force_cpu else 'cpu'
-
- # 设置SuperPoint和SuperGlue的配置参数
- config = {
- 'superpoint': {
- 'nms_radius': 4,
- 'keypoint_threshold': 0.005,
- 'max_keypoints': -1
- },
- 'superglue': {
- 'weights': opt.superglue,
- 'sinkhorn_iterations': 20,
- 'match_threshold': 0.2,
- }
- }
-
- # 创建Matching类的实例,用于图像匹配
- matching = Matching(config).eval().to(device)
- keys = ['keypoints', 'scores', 'descriptors']
-
- # 函数:处理图像尺寸调整
- def process_resize(w, h, resize):
- # 确保resize参数是合法的
- assert(len(resize) > 0 and len(resize) <= 2)
- # 如果只提供了一个值,基于最大维度调整比例
- if len(resize) == 1 and resize[0] > -1:
- scale = resize[0] / max(h, w)
- w_new, h_new = int(round(w*scale)), int(round(h*scale))
- # 如果提供的值是-1,保持原有尺寸
- elif len(resize) == 1 and resize[0] == -1:
- w_new, h_new = w, h
- else: # len(resize) == 2: # 如果提供了两个值,直接使用这两个值作为新的宽和高
- w_new, h_new = resize[0], resize[1]
-
- # 如果新的分辨率太小或太大,给出警告
- if max(w_new, h_new) < 160:
- print('警告:输入分辨率非常小,结果可能会有很大差异')
- elif max(w_new, h_new) > 2000:
- print('警告:输入分辨率非常大,可能会导致内存不足')
-
- return w_new, h_new
-
- # 定义load_image函数,用于加载和预处理图像
- def load_image(impath, resize):
- grayim = cv2.imread(impath, 0)
- # 以灰度模式读取图像
- if grayim is None:
- raise Exception('Error reading image %s' % impath)
- w, h = grayim.shape[1], grayim.shape[0]
- w_new, h_new = process_resize(w, h, resize)
- # 调用process_resize函数计算调整后的尺寸
- grayim = cv2.resize(grayim, (w_new, h_new), interpolation=cv2.INTER_AREA)
- # 使用cv2.resize函数调整图像尺寸
- return grayim
- # 返回调整后的灰度图像
-
-
- image_path_0 = "/home/fairlee/786D6A341753F4B4/KITTI/sequences_kitti_00_21/01/image_0/000000.png"
- frame0 = load_image(image_path_0, opt.resize)
-
- image_path_1 = "/home/fairlee/786D6A341753F4B4/KITTI/sequences_kitti_00_21/01/image_0/000001.png"
- frame1 = load_image(image_path_1, opt.resize)
-
-
-
- if __name__ == '__main__':
- # 将第一帧图像转换为张量,并移动到指定设备上
- frame_tensor0 = frame2tensor(frame0, device)
- # 使用SuperPoint提取第一帧图像的关键点和描述符
- last_data = matching.superpoint({'image': frame_tensor0})
- # 为第一帧图像的关键点、得分和描述符添加'0'后缀,以区分不同帧
- last_data = {k + '0': last_data[k] for k in keys}
- # 将第一帧图像的张量存储在last_data字典中
- last_data['image0'] = frame_tensor0
- # 存储第一帧图像
- last_frame = frame0
- # 存储第一帧图像的ID
- last_image_id = 0
-
- # 将第二帧图像转换为张量,并移动到指定设备上
- frame_tensor1 = frame2tensor(frame1, device)
- # 使用SuperGlue进行特征匹配,将第一帧图像的数据与第二帧图像的张量传递给matching函数
- pred = matching({**last_data, 'image1': frame_tensor1})
- # 获取第一帧图像的关键点坐标,并将其转换为NumPy数组
- kpts0 = last_data['keypoints0'][0].cpu().numpy()
- # 获取第二帧图像的关键点坐标,并将其转换为NumPy数组
- kpts1 = pred['keypoints1'][0].cpu().numpy()
-
- # 获取匹配结果,将其转换为NumPy数组
- matches = pred['matches0'][0].cpu().numpy()
- # 获取匹配置信度,将其转换为NumPy数组
- confidence = pred['matching_scores0'][0].cpu().numpy()
-
- # 找到有效的匹配,即匹配索引大于-1的位置
- valid = matches > -1
- # 获取第一帧图像中有效匹配的关键点坐标
- mkpts0 = kpts0[valid]
- # 获取第二帧图像中与第一帧图像有效匹配的关键点坐标
- mkpts1 = kpts1[matches[valid]]
-
- stem0, stem1 = last_image_id, 1
- # 打印匹配的关键点信息
- print(f"Matched keypoints in frame {stem0} and {stem1}:")
- for i, (kp0, kp1) in enumerate(zip(mkpts0, mkpts1)):
- print(f"Match {i}: ({kp0[0]:.2f}, {kp0[1]:.2f}) -> ({kp1[0]:.2f}, {kp1[1]:.2f})")
-
- 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, frame1, 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)
- cv2.waitKey(0)
- cv2.destroyAllWindows()
-
- if 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)
第二个版本的代码:
- #! /usr/bin/env python3
- import cv2
- import torch
- from models.matching import Matching
- from models.utils import (frame2tensor)
- torch.set_grad_enabled(False)
-
- # 设置SuperPoint和SuperGlue的配置参数
- config = {
- 'superpoint': {
- 'nms_radius': 4,
- 'keypoint_threshold': 0.005,
- 'max_keypoints': -1
- },
- 'superglue': {
- 'weights': 'outdoor',
- 'sinkhorn_iterations': 20,
- 'match_threshold': 0.2,
- }
- }
-
- device = 'cuda' if torch.cuda.is_available() else 'cpu'
- # 创建Matching类的实例,用于图像匹配
- matching = Matching(config).eval().to(device)
- keys = ['keypoints', 'scores', 'descriptors']
-
- # 对于灰度图像,返回的NumPy数组将是一个二维数组,其中数组的形状对应于图像的高度和宽度(H x W)。
- # 每个元素的值代表了对应像素的亮度,通常是一个0到255的整数(对于8位灰度图像)。
- frame0 = cv2.imread("/home/fairlee/000001.jpg", 0)
- frame1 = cv2.imread("/home/fairlee/000000.jpg", 0)
-
- def match_frames(frame0, frame1, device, matching, keys):
- """
- Match keypoints between two frames and return the matched coordinates and confidence scores.
- Parameters:
- - frame0: Numpy array, first image frame.
- - frame1: Numpy array, second image frame.
- - device: The device to perform computation on.
- - matching: Matching object with a method to match points between frames.
- - keys: List of keys to extract from the matching data.
- Returns:
- A tuple of (mkpts0, mkpts1, confidence_scores), where:
- - mkpts0: Matched keypoints in the first frame.
- - mkpts1: Matched keypoints in the second frame.
- - confidence_scores: Confidence scores of the matches.
- """
- # Convert frames to tensors and move to the device
- frame_tensor0 = frame2tensor(frame0, device)
- frame_tensor1 = frame2tensor(frame1, device)
-
- # Get data from the first frame
- last_data = matching.superpoint({'image': frame_tensor0})
- last_data = {k + '0': last_data[k] for k in keys}
- last_data['image0'] = frame_tensor0
-
- # Perform matching
- pred = matching({**last_data, 'image1': frame_tensor1})
-
- # Extract keypoints and convert to Numpy arrays
- kpts0 = last_data['keypoints0'][0].cpu().numpy()
- kpts1 = pred['keypoints1'][0].cpu().numpy()
-
- # Extract matches and confidence scores, convert to Numpy arrays
- matches = pred['matches0'][0].cpu().numpy()
- confidence = pred['matching_scores0'][0].cpu().numpy()
-
- # Filter valid matches
- valid = matches > -1
- mkpts0 = kpts0[valid]
- mkpts1 = kpts1[matches[valid]]
-
- return mkpts0, mkpts1, confidence[valid]
结果:
通过运行这段代码,我们可以看到SuperGlue算法在图像特征匹配方面的强大能力。代码首先处理输入图像,然后使用SuperPoint模型提取特征点和描述子,接着SuperGlue模型根据描述子进行关键点匹配。匹配过程的结果会被可视化显示出来,如果指定了输出目录,还会将结果图像保存下来。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。