SuperGluePretrainedNetwork调用接口版本(两个版本!)

        本脚本是一个基于Python的应用,旨在演示如何使用SuperGlue算法进行图像之间的特征匹配。SuperGlue是一个强大的特征匹配工具,能够在不同的图像之间找到对应的关键点。这个工具尤其适用于计算机视觉任务,如立体视觉、图像拼接、对象识别和追踪等场景。脚本使用PyTorch框架,并且可以选择在CPU或GPU上运行。

脚本的工作流程如下:

  1. 解析命令行参数,用于设置输入输出目录、图像尺寸、SuperGlue配置等。
  2. 根据用户选择,决定算法是在CPU还是GPU上执行。
  3. 加载预设的配置,初始化SuperPoint和SuperGlue模型。
  4. 定义图像预处理函数来调整图像大小。
  5. 加载两幅图像,调整它们的大小,并将它们转换为PyTorch张量。
  6. 使用SuperPoint提取关键点和描述符。
  7. 使用SuperGlue算法匹配两幅图像的关键点。
  8. 可视化并打印匹配的关键点坐标。
  9. 如果设置了输出目录,将结果图像写到磁盘上。

        这个脚本展示了如何在实践中使用深度学习模型来处理实际问题,并提供了图像匹配演示。

#! /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模型根据描述子进行关键点匹配。匹配过程的结果会被可视化显示出来,如果指定了输出目录,还会将结果图像保存下来。

  • 3
    点赞
  • 10
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值