D2-Net：用于局部特征联合检测与描述的可训练 CNN

本仓库包含以下论文的实现：

“D2-Net：用于局部特征联合检测与描述的可训练 CNN”。
M. Dusmanu、I. Rocco、T. Pajdla、M. Pollefeys、J. Sivic、A. Torii 和 T. Sattler。CVPR 2019。

论文 arXiv 链接，项目页面

快速入门

建议使用 Python 3.6 及以上版本运行我们的代码。可以使用 Conda 安装所需的软件包：

conda install pytorch torchvision cudatoolkit=10.0 -c pytorch
conda install h5py imageio imagesize matplotlib numpy scipy tqdm

下载模型

现成的 Caffe VGG16 权重及其调优版本可以通过以下命令下载：

mkdir models
wget https://dusmanu.com/files/d2-net/d2_ots.pth -O models/d2_ots.pth
wget https://dusmanu.com/files/d2-net/d2_tf.pth -O models/d2_tf.pth
wget https://dusmanu.com/files/d2-net/d2_tf_no_phototourism.pth -O models/d2_tf_no_phototourism.pth

更新 - 2019年5月23日 我们新增了一组在 MegaDepth 数据集上训练的权重，该数据集去除了 PhotoTourism 场景（sagrada_familia - 0019、lincoln_memorial_statue - 0021、british_museum - 0024、london_bridge - 0025、us_capitol - 0078、mount_rushmore - 1589）。初步结果表明性能相近。要在测试时使用这些权重，需添加 --model_file models/d2_tf_no_phototourism.pth。

特征提取

extract_features.py 可用于从给定的图像列表中提取 D2 特征。单尺度特征处理 1200×1600 分辨率的图像时所需显存不到 6GB。使用 --multiscale 标志可提取多尺度特征；此时建议至少配备 12GB 显存。

输出格式可以是 npz 或 mat。无论哪种格式，特征文件都包含三个数组：

• keypoints [N x 3] 数组，存储关键点的位置 x, y 和尺度 s。位置遵循 COLMAP 格式，其中 X 轴指向右侧，Y 轴指向下方。
• scores [N] 数组，存储关键点的激活值（值越高越好）。
• descriptors [N x 512] 数组，存储 L2 归一化的描述子。

python extract_features.py --image_list_file images.txt (--multiscale)

使用 kapture 数据集进行特征提取

Kapture 是一种基于文本和二进制文件的枢纽文件格式，用于描述 SFM（运动恢复结构）以及更广泛的传感器采集数据。

其 GitHub 地址为：https://github.com/naver/kapture。该库包含常用格式的转换工具，并且多个流行的数据集可以直接以 kapture 格式提供。

可通过以下命令安装：

pip install kapture

数据集可按如下方式下载：

kapture_download_dataset.py update
kapture_download_dataset.py list
# 例如：安装 Extended-CMU-Seasons_slice22 的映射和查询数据
kapture_download_dataset.py install "Extended-CMU-Seasons_slice22_*"

若要将您自己的数据集转换为 kapture 格式，请参阅 此处 的示例。

安装完成后，您可以使用以下命令为 kapture 数据集提取关键点：

python extract_kapture.py --kapture-root pathto/yourkapturedataset (--multiscale)

运行 python extract_kapture.py --help 以获取更多关于提取参数的信息。

在 MegaDepth 上调优

此处提供的训练流程是 TensorFlow 代码的 PyTorch 实现，该 TensorFlow 代码曾用于训练并可在此处下载的模型。

更新 - 2019年6月5日 我们修复了数据预处理中的一个错误；重新训练现在能得到与原始 TensorFlow 实现相似的结果。

更新 - 2019年8月7日 我们发布了更新且更精确的训练数据集版本；在达到相同性能的情况下，训练更加稳定且速度显著提升。

下载并预处理 MegaDepth 数据集

为此部分，应先安装 COLMAP。请参考其官方网站获取安装说明。

下载完整的 MegaDepth 数据集（包括 SfM 模型）后，第一步是生成去畸变后的重建结果。这可以通过调用 undistort_reconstructions.py 来完成：

python undistort_reconstructions.py --colmap_path /path/to/colmap/executable --base_path /path/to/megadepth

接下来，可以使用 preprocess_megadepth.sh 获取相机参数，并计算所有场景中图像之间的重叠区域。

bash preprocess_undistorted_megadepth.sh /path/to/megadepth /path/to/output/folder

训练

在下载并预处理 MegaDepth 数据集后，即可立即开始训练：

python train.py --use_validation --dataset_path /path/to/megadepth --scene_info_path /path/to/preprocessing/output

BibTeX

如果您在项目中使用此代码，请引用以下论文：

@InProceedings{Dusmanu2019CVPR,
    author = {Dusmanu, Mihai and Rocco, Ignacio and Pajdla, Tomas and Pollefeys, Marc and Sivic, Josef and Torii, Akihiko and Sattler, Torsten},
    title = {{D2-Net: A Trainable CNN for Joint Detection and Description of Local Features}},
    booktitle = {2019 IEEE/CVF 计算机视觉与模式识别会议论文集},
    year = {2019},
}

D2-Net 快速上手指南

D2-Net 是一个可训练的卷积神经网络（CNN），用于联合检测和描述局部特征。本指南将帮助你快速配置环境、下载模型并提取图像特征。

1. 环境准备

• 操作系统：Linux / macOS (Windows 需自行配置 CUDA 环境)
• Python 版本：推荐 Python 3.6 及以上
• 硬件要求：
  • 单尺度特征提取：显存至少 6GB (针对 1200x1600 分辨率图像)
  • 多尺度特征提取：推荐显存 12GB 及以上
• 前置依赖：PyTorch, CUDA Toolkit

2. 安装步骤

2.1 创建 Conda 环境并安装依赖

推荐使用 Conda 管理环境。以下命令将安装 PyTorch 及项目所需的所有 Python 包。

提示：国内用户建议使用清华或中科大镜像源加速下载。

# 创建新环境 (可选)
conda create -n d2net python=3.8
conda activate d2net

# 安装 PyTorch 和 CUDA 工具包 (根据实际 CUDA 版本调整，此处以原文 10.0 为例)
# 国内加速示例 (清华源):
conda install pytorch torchvision cudatoolkit=10.0 -c pytorch -c https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/

# 安装其他依赖包
conda install h5py imageio imagesize matplotlib numpy scipy tqdm
# 若 conda 安装慢，也可使用 pip + 国内源:
# pip install h5py imageio imagesize matplotlib numpy scipy tqdm -i https://pypi.tuna.tsinghua.edu.cn/simple

2.2 下载预训练模型

创建模型目录并下载官方提供的权重文件（包含 Caffe VGG16 基础权重及微调后的权重）。

mkdir models
wget https://dusmanu.com/files/d2-net/d2_ots.pth -O models/d2_ots.pth
wget https://dusmanu.com/files/d2-net/d2_tf.pth -O models/d2_tf.pth
wget https://dusmanu.com/files/d2-net/d2_tf_no_phototourism.pth -O models/d2_tf_no_phototourism.pth

注：d2_tf_no_phototourism.pth 是在 MegaDepth 数据集上训练且排除了 PhotoTourism 场景的权重，通常在测试时表现更佳。

3. 基本使用

3.1 准备图像列表

首先，创建一个文本文件（例如 images.txt），每行包含一张待处理图像的绝对路径或相对路径。

/path/to/image1.jpg
/path/to/image2.jpg

3.2 提取特征

运行 extract_features.py 脚本提取特征。默认输出为单尺度特征，若需更高精度可添加 --multiscale 参数（需更大显存）。

单尺度提取（推荐入门）：

python extract_features.py --image_list_file images.txt

多尺度提取（高精度，需 >=12GB 显存）：

python extract_features.py --image_list_file images.txt --multiscale

3.3 输出结果说明

脚本执行完成后，将生成 .npz 或 .mat 格式的特征文件。每个文件包含以下三个数组：

• keypoints: [N x 3] 数组，包含关键点坐标 x, y 和尺度 s（遵循 COLMAP 格式，X 轴向右，Y 轴向下）。
• scores: [N] 数组，关键点的激活分数（值越高越好）。
• descriptors: [N x 512] 数组，L2 归一化后的特征描述子。

---

如需处理 Kapture 格式数据集，请先安装 pip install kapture，然后使用 python extract_kapture.py --kapture-root <数据集路径> 进行提取。