SiLK - 简单可学习的关键点

[Arxiv论文]
作者：Pierre Gleize、Weiyao Wang 和 Matt Feiszli
会议：ICCV 2023

SiLK 是一个用于学习关键点的自监督框架。SiLK 注重 简单性 和 灵活性，同时在现有基准测试上提供了最先进的、具有竞争力的结果。

还提供了预训练模型。

发布的代码已在 Linux 系统上进行了测试，使用两块 Tesla V100-SXM2 GPU 进行训练，大约需要 5 小时。

需求

• 必须安装 conda 才能设置 silk 环境。
• 训练 SiLK 需要 两 块 GPU。

使用方法

• 如何设置 Python 环境？
• 如何准备数据集？
• 如何训练 SiLK？
• 如何添加主干网络？
• 如何运行评估流程？
• 如何进行推理？
• 如何将 SiLK 转换为 Torch Script？
• 如何将 SiLK 特征导入 COLMAP？
• 常见问题解答

结果

以下结果是使用我们的 VGG-4 主干网络计算得出的（表 2、3、6 使用检查点 pvgg-4.ckpt，表 4、5 使用 coco-rgb-aug）。为了复现 IMC2022 的结果，我们还提供了 Kaggle 笔记本。

完整文档

我们提供了一份文档，但并不全面。如果对代码的某些部分需要进一步说明，请提交一个新的问题。我们将根据社区的需求补充文档。

我们的文档可以在这里找到：doc/silk/index.html。

贡献

请参阅 CONTRIBUTING 文件，了解如何参与贡献。

许可证

SiLK 根据 LICENSE 文件中的规定，采用 GNU 通用公共许可证（第 3 版） 许可。

SiLK 快速上手指南

SiLK (Simple Learned Keypoints) 是一个用于学习关键点的自监督框架，以其简洁性和灵活性著称，并在多个基准测试中取得了最先进的结果。本指南将帮助中国开发者快速搭建环境并运行模型。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (代码已在 Linux 环境下测试通过)
• 硬件要求:
  • 训练: 需要 2 块 GPU (原文测试环境为 Tesla V100-SXM2)。
  • 推理: 单卡即可运行。
• 软件依赖:
  • Conda (必须安装，用于管理 Python 环境)
  • CUDA toolkit (需与您的 GPU 驱动及 PyTorch 版本兼容)

提示: 国内用户建议使用清华源或中科大源加速 Conda 和 Pip 包的下载。

安装步骤

1. 创建并激活 Conda 环境

首先克隆项目代码（假设您已获取仓库），然后创建专用环境：

git clone https://github.com/facebookresearch/silk.git
cd silk

# 创建名为 silk 的 conda 环境 (具体 python 版本请参考项目 requirements 或 setup 文档，通常为 3.8+)
conda create -n silk python=3.8
conda activate silk

2. 安装依赖包

进入项目目录后，按照官方提供的设置脚本安装依赖。如果 setup.md 中有具体的 pip install 命令，请优先执行。通常步骤如下：

# 推荐使用国内镜像源加速安装
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

# 如果项目提供特定的安装脚本，请执行：
# bash doc/usage/setup.sh (如有)

注意：由于训练需要多卡支持，请确保安装的 torch 和 torchvision 版本支持您的 CUDA 版本。

3. 下载预训练模型

为了快速验证，您可以直接下载官方提供的预训练权重：

• VGG-4 backbone (用于部分基准测试): pvgg-4.ckpt
• COCO 增强版 (用于其他基准测试): coco-rgb-aug.ckpt

将下载的 .ckpt 文件放置在项目指定的 assets/models 目录下（具体路径可参考 doc/usage/inference.md）。

基本使用

以下是最简单的推理示例，展示如何加载模型并提取关键点特征。

运行推理 (Inference)

假设您已准备好两张图像用于匹配，可以使用以下命令运行推理脚本（具体脚本名称请以仓库实际结构为准，通常为 python scripts/inference.py 或类似）：

# 示例命令：运行推理并保存匹配结果
# 请根据实际脚本路径调整参数
python -m silk.inference \
    --checkpoint_path ./assets/models/silk/coco-rgb-aug.ckpt \
    --image_pair ./data/images/img1.png ./data/images/img2.png \
    --output_dir ./results

导出特征到 COLMAP

如果您需要将 SiLK 提取的特征用于三维重建（COLMAP），可以使用专门的转换工具：

# 将 SiLK 特征转换为 COLMAP 格式
python -m silk.utils.colmap_export \
    --input_dir ./images \
    --checkpoint_path ./assets/models/silk/coco-rgb-aug.ckpt \
    --output_dir ./colmap_output

简单训练示例

若您拥有双卡环境并希望从头训练（约需 5 小时）：

# 启动分布式训练 (需要 2 张 GPU)
python -m torch.distributed.launch --nproc_per_node=2 train.py \
    --config configs/silk_default.yaml \
    --data_root ./datasets/megadepth

注意: 详细的数据集准备 (dataset.md) 和配置文件修改请参考项目 doc/usage/ 目录下的详细文档。