Stereo-Waterdrop-Removal
This repository includes official codes for "Stereo Waterdrop Removal with Row-wise Dilated Attention (IROS2021)".
Stereo Waterdrop Removal with Row-wise Dilated Attention
Zifan Shi, Na Fan, Dit-Yan Yeung, Qifeng Chen
HKUST
Existing vision systems for autonomous driving or robots are sensitive to waterdrops adhered to windows or camera lenses. Most recent waterdrop removal approaches take a single image as input and often fail to recover the missing content behind waterdrops faithfully. Thus, we propose a learning-based model for waterdrop removal with stereo images. A real-world dataset that contains stereo images with and without waterdrops is provided to benefit the related research.
Installation
Clone this repo.
git clone https://github.com/VivianSZF/Stereo-Waterdrop-Removal.git
cd Stereo-Waterdrop-Removal/
We have tested our code on Ubuntu 18.04 LTS with PyTorch 1.6.0 and CUDA 10.2. Please install dependencies by
conda env create -f environment.yml
Datasets
The dataset can be downloaded from the link.
'train', 'val' and 'test' refer to training set, validation set and test set captured by ZED 2. 'test_mynt' contains test images from MYNT EYE camera. In each folder, '000' denotes the waterdrop-free image (Ground truth). 'xxx_0' is the left image while 'xxx_1' is the right image. The dataset can be put under the 'dataset' folder.
Training
The arguments for training are listed in train.py. To train the model, run with the following code
sh train.sh
The checkpoints and the validation ressults will be saved into ./result/{exp_name}/train/.
Test
Download the pretrained checkpoints and put them under ./result/{exp_name}/train/. The arguments for test are listed in test.py. You can specify them in test.sh and run the command
sh test.sh
The output images are available under ./result/{exp_name}/test/
Citation
@inproceedings{shi2021stereo,
title = {Stereo Waterdrop Removal with Row-wise Dilated Attention},
author = {Shi, Zifan and Fan, Na and Yeung, Dit-Yan and Chen, Qifeng},
booktitle = {IROS},
year = {2021}
}
