<!-- WEHUB_ZH_README -->
> [!NOTE]
> 本文档由 WeHub 基于上游 README 翻译整理，属于社区翻译，非官方中文文档。  
> [English](./README.en.md) · [原始项目](https://github.com/megvii-research/CREStereo) · [上游 README](https://github.com/megvii-research/CREStereo/blob/HEAD/README.md)  
> 原作者、版权与许可证归属以原始项目及本仓库 LICENSE 文件为准。

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


本仓库包含我们论文的 [MegEngine](https://github.com/MegEngine/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)**

**[论文](https://openaccess.thecvf.com/content/CVPR2022/papers/Li_Practical_Stereo_Matching_via_Cascaded_Recurrent_Network_With_Adaptive_Correlation_CVPR_2022_paper.pdf)  | [ArXiv](https://arxiv.org/abs/2203.11483) | [BibTeX](#citation)** 

<img src="img/teaser.jpg">

## 数据集

### 本文提出的数据集

#### 下载

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

- 使用 shell 脚本下载 `dataset_download.sh`

```shell
sh dataset_download.sh
```

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

 - 从百度云 [此处](https://pan.baidu.com/s/1iB96-ftCgPFTlrj220qw8Q)(Extraction code: `aa3g`) 下载，并手动解压 tar 文件。

#### 视差格式

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

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

### 其他公开数据集

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

 - [SceneFlow](https://lmb.informatik.uni-freiburg.de/resources/datasets/SceneFlowDatasets.en.html)
 - [Sintel](http://sintel.is.tue.mpg.de/stereo)
 - [Middlebury](https://vision.middlebury.edu/stereo/data/)
 - [ETH3D](https://www.eth3d.net/datasets#low-res-two-view-training-data)
 - [KITTI 2012/2015](http://www.cvlibs.net/datasets/kitti/eval_stereo.php) 
 - [Falling Things](https://research.nvidia.com/publication/2018-06_Falling-Things)
 - [InStereo2K](https://github.com/YuhuaXu/StereoDataset)
 - [HR-VS](https://drive.google.com/file/d/1SgEIrH_IQTKJOToUwR1rx4-237sThUqX/view)


 ## 依赖


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

```bash
python3 -m pip install -r requirements.txt
```

我们还提供 Docker，以便快速运行代码：
```bash
docker run --gpus all -it -v /tmp:/tmp ylmegvii/crestereo
shotwell /tmp/disparity.png
```

 ## 推理

从 [此处](https://drive.google.com/file/d/1Wx_-zDQh7BUFBmN9im_26DFpnf3AkXj4/view) 下载预训练的 MegEngine 模型并运行：

```shell
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` 中修改配置，并运行以下命令：

```shell
python3 train.py
```

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

```shell
tensorboard --logdir ./train_log
```

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


## 致谢
部分代码改编自先前的工作：
 - [RAFT](https://github.com/princeton-vl/RAFT)(code base)
 - [LoFTR](https://github.com/zju3dv/LoFTR)(attention module)
 - [HSMNet](https://github.com/gengshan-y/high-res-stereo)(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}
}
```