DerainNeRF: 3D Scene Estimation with Adhesive Waterdrop Removal

Yunhao Li, Jing Wu, Lingzhe Zhao and Peidong Liu [ICRA 2024]

This is a simplified implementation of the paper DerainNeRF: 3D Scene Estimation with Adhesive Waterdrop Removal (ICRA 2024). DerainNeRF reconstruct the clear 3D scene implicitly from multi-view images degraded by adhesive waterdrops.

✨News

📺 [2024.01] Our paper has been accepted by IEEE ICRA 2024!

Novel View Synthesis

Image Waterdrop Removal Result

Method overview

Our method exploits a pre-trained attention network to predict the location of waterdrops and then train a Neural Radiance Fields (NeRF) to recover the 3D scene implicitly.

Quickstart

1. Setup environment

git clone https://github.com/yunhaoli2020/DerainNeRF
cd DerainNeRF
pip install -r requirements.txt

2. Download datasets

You can download the data here.

For different scenes (cozy2room, factory etc.), the folder images includes images in original resolution and the folder images_3 (or images_4) includes input images of DerainNeRF. The number indicates the scale factor. For example, images_3 indicates the 3x downsample when loading the dataset. The file poses_bounds.npy contains poses derived from COLMAP. The file rainMask.npy contains binary masks which indicates the waterdrops in input images.

3. Configs

Change the data path and other parameters (if needed) in configs/tanabata.txt. We use tanabata scene as an example.

4. Training

python run_derainnerf.py --config configs/tanabata.txt

After training, you can get clear images, depth maps and novel-view images synthesized from NeRF.

5. Demo with our pre-trained model

You can test our code and render clear images with the provided weight files. To do this, you should first download the pre-trained models from here, unzip it, then put the weight file under the corresponding logs folder ./logs, finally run

python run_derainnerf.py --config configs/tanabata.txt

Your own data

images: This folder is used to estimate initial camera poses from waterdrop images. Specifically, just put your own data in the folder images, and run imgs2poses.py script from the LLFF code to estimate camera poses and generate poses_bounds.npy.

images_X: This is the default training folder, which includes the downsampled waterdrop images in images folder. The number X indicates the downsample rate. If X it 1, then the images in the two folders would be the same.

rainMask.npy: This file contains binary masks which indicates waterdrops in the input images.


images folder: (suppose 20 images)
000.png
001.png
...
019.png
#-----------------------------------------------------------------------------------------#
images_X folder: (training folder, suppose 20 images)
000.png
001.png
....
019.png
#-----------------------------------------------------------------------------------------#
poses_bounds.npy
rainMask.npy

Citation

If you find this useful, please consider citing our paper:

@misc{li2024derainnerf,
      title={DerainNeRF: 3D Scene Estimation with Adhesive Waterdrop Removal}, 
      author={Yunhao Li and Jing Wu and Lingzhe Zhao and Peidong Liu},
      year={2024},
      eprint={2403.20013},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

Acknowledgment

The overall framework, metrics computing and camera transformation are derived from nerf-pytorch and AttGAN respectively. We appreciate the effort of the contributors to these repositories.

yunhaoli2020 / DerainNeRF