【AI】使用LoFTR进行图像匹配测试Demo

LoFTR图像匹配的源码解析我们在上篇文章中已经写了，对于怎么试用一下，我这边再啰嗦一下。

0.环境搭建

详细的搭建教程请点击链接查看，这里只对需要特殊注意的地方做阐述
1.创建的Python环境采用python3.8的环境，因为文章发布较早，使用3.10版本测试的时候发现了依赖版本不兼容的问题
2.torch试用whl文件的安装方式比较稳妥，然后在源码requirements.txt中将torch的依赖注释掉，在安装完torch之后使用pip install -r requirements.txt进行安装依赖。

1.运行文件准备

1.下载预训练权重文件，由于我这边使用的是外景图片，所以下载了outdoor的模型文件。
官方提供了下载地址：https://drive.google.com/drive/folders/1DOcOPZb3-5cWxLqn256AhwUVjBPifhuf，需要一点魔法
2.图片文件随便两张带有重叠景象的图片，我这边随便用了两张
3.注意修改demo文件中的权重路径和图片路径

2.运行demo文件

这里我是在模型源码的demo文件夹下面进行的测试，如果在其他路径下需要注意引入路径问题。

import torch
import cv2
import numpy as np
import matplotlib.cm as cm

from src.utils.plotting import make_matching_figure
from src.loftr import LoFTR, default_cfg


if __name__ == '__main__':
    # 根据图片拍摄场景和下载的预训练模型进行选择 可选：indoor、outdoor
    image_type = 'outdoor'
	# 根据个人图片进行修改
    img0_pth = "images/12003890_f6c899bec0_o.jpg"
    img1_pth = "images/13866250_56e0509621_o.jpg"
    image_pair = [img0_pth, img1_pth]

    # The default config uses dual-softmax.
    # The outdoor and indoor models share the same config.
    # You can change the default values like thr and coarse_match_type.
    matcher = LoFTR(config=default_cfg)
    # load中修改预训练模型地址
    if image_type == 'indoor':
        matcher.load_state_dict(torch.load("indoor_ds.ckpt")['state_dict'])
    elif image_type == 'outdoor':
        matcher.load_state_dict(torch.load("outdoor_ds.ckpt")['state_dict'])
    else:
        raise ValueError("Wrong image_type is given.")
    matcher = matcher.eval().cuda()

    # Rerun this cell (and below) if a new image pair is uploaded.
    img0_raw = cv2.imread(image_pair[0], cv2.IMREAD_GRAYSCALE)
    img1_raw = cv2.imread(image_pair[1], cv2.IMREAD_GRAYSCALE)
    img0_raw = cv2.resize(img0_raw, (640, 480))
    img1_raw = cv2.resize(img1_raw, (640, 480))

    img0 = torch.from_numpy(img0_raw)[None][None].cuda() / 255.
    img1 = torch.from_numpy(img1_raw)[None][None].cuda() / 255.
    batch = {'image0': img0, 'image1': img1}

    # Inference with LoFTR and get prediction
    with torch.no_grad():
        matcher(batch)
        mkpts0 = batch['mkpts0_f'].cpu().numpy()
        mkpts1 = batch['mkpts1_f'].cpu().numpy()
        mconf = batch['mconf'].cpu().numpy()

    # Draw
    color = cm.jet(mconf, alpha=0.7)
    text = [
        'LoFTR',
        'Matches: {}'.format(len(mkpts0)),
    ]
    fig = make_matching_figure(img0_raw, img1_raw, mkpts0, mkpts1, color, mkpts0, mkpts1, text)

    # A high-res PDF will also be downloaded automatically.
    make_matching_figure(img0_raw, img1_raw, mkpts0, mkpts1, color, mkpts0, mkpts1, text, path="LoFTR-colab-demo.pdf")