tyz1030 / neuralsea

[RAL+ICRA 2024] Underwater color correction is now True and Dense! Thanks to NeRF. No color assumptions, no pre-train, no depth estimation needed. Purely physics-based. Built for Robots!

Beyond NeRF Underwater: Learning Neural Reflectance Fields for True Color Correction of Marine Imagery

[UPDATE!] Our paper just gets accepted by RA-L! Thanks to reviewers, collaborators, editors, and those who raise issues on GitHub which helped us improve the code a lot.

All data presented in the paper are now released.
Example: Download 'data_lego_white' and put it under the neuralsea directory. Download white lego weights and put it under neuralsea/checkpoints/ to use our pretrained weights (not all weights are released so far. I will upload them once I got time. Feel free to train from scratch and I don't usually stick with certain random seeds.).

Publication

Our paper is published at IEEE RA-L. You can also find it on arxiv. To appear on ICRA 2024.
This work is supported by National Oceanic and Atmospheric Administration (NOAA) under grant NA22OAR0110624.

Visualizations

Left: with water effects; Right: color corrected

Groundtruth image

Groundtruth image

More real-world results:
Lake Erie:


Water Tank Low Turbidity:

Water Tank Mid Turbidity:

Water Tank High Turbidity:

Dependencies

Install PyTorch:

pip install torch torchvision

Install PyTorch3D, please follow their instruction. We use the following to install:

pip install "git+https://github.com/facebookresearch/pytorch3d.git"

Install other dependencies:

pip install hydra-core plotly visdom matplotlib

Train & Test

# for synthetic data example
python3 train_nerf.py --config-name synthetic_lego_white
python3 test_nerf.py --config-name synthetic_lego_white

# for water tank data example
python3 train_nerf.py --config-name real_watertank
python3 test_nerf.py --config-name real_watertank

Visualization

install visdom

pip install visdom

run visdom

visdom

Then in your browser, navigate to http://localhost:8097/

Paper

If you find this study helpful please kindly cite us:

@ARTICLE{10225666,
  author={Zhang, Tianyi and Johnson-Roberson, Matthew},
  journal={IEEE Robotics and Automation Letters}, 
  title={Beyond NeRF Underwater: Learning Neural Reflectance Fields for True Color Correction of Marine Imagery}, 
  year={2023},
  volume={8},
  number={10},
  pages={6467-6474},
  doi={10.1109/LRA.2023.3307287}}

Miscellaneous

(Irrelavant to this paper) Our field work featured on noaa.gov