megvii-research--crestereo/README.md

Note

本文档由 WeHub 基于上游 README 翻译整理，属于社区翻译，非官方中文文档。
English · 原始项目 · 上游 README
原作者、版权与许可证归属以原始项目及本仓库 LICENSE 文件为准。

[CVPR 2022] Practical Stereo Matching via Cascaded Recurrent Network with Adaptive Correlation

本仓库包含我们论文的 MegEngine 实现：

Practical Stereo Matching via Cascaded Recurrent Network with Adaptive Correlation
Jiankun Li, Peisen Wang, Pengfei Xiong, Tao Cai, Ziwei Yan, Lei Yang, Jiangyu Liu, Haoqiang Fan, Shuaicheng Liu
CVPR 2022 (Oral)

论文 | ArXiv | BibTeX

数据集

本文提出的数据集

下载

下载我们论文中提出的数据集（约 400GB）有 两种方式：

• 使用 shell 脚本下载 dataset_download.sh

sh dataset_download.sh

数据集将被下载并解压到 ./stereo_trainset/crestereo

• 从百度云 此处(Extraction code: aa3g) 下载，并手动解压 tar 文件。

视差格式

视差以 .png uint16 格式保存，可使用 opencv 的 imread 函数加载：

def get_disp(disp_path):
    disp = cv2.imread(disp_path, cv2.IMREAD_UNCHANGED)
    return disp.astype(np.float32) / 32

其他公开数据集

我们使用的其他公开数据集包括

• SceneFlow
• Sintel
• Middlebury
• ETH3D
• KITTI 2012/2015
• Falling Things
• InStereo2K
• HR-VS

依赖

CUDA 版本：10.1，Python 版本：3.6.9

• MegEngine v1.8.2
• opencv-python v3.4.0
• numpy v1.18.1
• Pillow v8.4.0
• tensorboardX v2.1

python3 -m pip install -r requirements.txt

我们还提供 Docker，以便快速运行代码：

docker run --gpus all -it -v /tmp:/tmp ylmegvii/crestereo
shotwell /tmp/disparity.png

推理

从 此处 下载预训练的 MegEngine 模型并运行：

python3 test.py --model_path path_to_mge_model --left img/test/left.png --right img/test/right.png --size 1024x1536 --output disparity.png

训练

在 cfgs/train.yaml 中修改配置，并运行以下命令：

python3 train.py

你可以启动 TensorBoard 来监控训练过程：

tensorboard --logdir ./train_log

并在浏览器中访问 http://localhost:6006 页面。

致谢

部分代码改编自先前的工作：

• RAFT(code base)
• LoFTR(attention module)
• HSMNet(data augmentaion)

感谢各位作者提供的优秀代码仓库。

引用

如果你在研究中发现本代码或数据集有帮助，请引用：

@inproceedings{li2022practical,
  title={Practical stereo matching via cascaded recurrent network with adaptive correlation},
  author={Li, Jiankun and Wang, Peisen and Xiong, Pengfei and Cai, Tao and Yan, Ziwei and Yang, Lei and Liu, Jiangyu and Fan, Haoqiang and Liu, Shuaicheng},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={16263--16272},
  year={2022}
}