This repository contains the code for AAAI 2024 paper IINet: Implicit Intra-inter Information Fusion for Real-Time Stereo Matching
paper-link
- Python 3.8
- CUDA 10.3
- PyTorch 1.11.0
matplotlib antialiased-cnns opencv-python progressbar2 timm==0.6.5 tensorboardX h5py kornia pypng
- SceneFlow Pre-trained Weights [google drive]
- Spring Pre-trained Weights [google drive]
Sceneflow dataset
Download the finalpass data of the Sceneflow dataset as well as the Disparity data.
Spring dataset
Download Spring dataset ,
and unzip train/test_frame_left/right.zip, train_disp1_left/right.zip.
The datasets are organized as follows
└── datasets
├── spring
| ├── test
| │ ├── 0003
| │ ├── 0019
| | :
| |
| ├── train
| │ ├── 0001
| │ └── 0002
| | :
|
└── sceneflow
├── disparity
│ ├── TEST
│ │ ├──A
│ │ ├──B
│ │ └──C
│ └── TRAIN
│ ├──15mm_focallength
│ :
│ ├──a_rain_of_stones_x2
│ :
│ └──A
├── frames_finalpass
├── TEST
│ ├──A
│ ├──B
│ └──C
└── TRAIN
├──15mm_focallength
:
Please download the pretrained sceneflow checkpoint and put it in the root directory. Then put left and right image in the 'test' folder and name them with 'left.png' and 'right.png'. Then run:
bash scripts/test_single.sh
The colored disparity map and 16bit disparity map will be saved in the same folder.
Set dataset_path in 'configs/data/sceneflow.yaml' and load_weights_fron_checkpoint in './scripts/eval_sceneflow.sh'. Then run:
bash scripts/eval_sceneflow.sh
bash scripts/eval_spring.sh
Set dataset_path in 'configs/data/sceneflow.yaml'. Then run:
bash scripts/train_sceneflow.sh
After training the model on Sceneflow, there will be saved checkpoints in './run/sceneflow/expname/models'. Set load_weights_from_checkpoint in './scripts/finetune_spring.sh'. Then run:
bash scripts/finetune_spring.sh
Run the script, and the gpu latency will be tested using torch.profiler. The screen outputs show the total latency of 3 forward process. Two files will be generated, torch_cuda_stack.json and torch_cpu_stack.json, which can be further loaded to speedscope to show latency of each module.
bash scripts/test_speed.sh
If you find our work useful in your research, please consider citing our paper
@inproceedings{li2024iinet,
title={IINet: Implicit Intra-inter Information Fusion for Real-Time Stereo Matching},
author={Li, Ximeng and Zhang, Chen and Su, Wanjuan and Tao, Wenbing},
booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
volume={38},
number={4},
pages={3225--3233},
year={2024}
}
Part of the code is adopted from previous works: Simplerecon, STTR