Skip to content
bilityniuPublic

Repository files navigation

3D-USE 3D-USE: From Image-Level to Scene-Level Underwater Enhancement

Jieyu Yuan1 · Yuanlin Zhang1 · Jihong Li1 · Chunle Guo1,2 · Huimin Lu3,4 · Chongyi Li1,2,*

1Nankai University    2NKIARI, Shenzhen Futian
3Southeast University    4Advanced Ocean Institute of Southeast University, Nantong
*Corresponding author

arXiv Project Page Code


📰 News

  • [2026.09.01] Code repository is now available! [Code]
  • [2026.08.28] Paper released on arXiv! [arXiv]

🌊 Overview

3D-USE aims to reconstruct an underwater scene with improved visibility and consistent appearance across novel viewpoints. Standard reconstruction preserves underwater degradation, while independent 2D enhancement can introduce view-dependent corrections. Our two-stage framework first builds a medium-aware Gaussian scene and then transfers paired 2D enhancement knowledge into a persistent 3D appearance representation, enabling enhanced novel-view rendering without a 2D UIE model at inference.

  • Stage 1: Medium-aware Gaussian reconstruction with MediumRBF.
  • Stage 2: ATC-guided scene enhancement stored in U-BAF.
  • Inference: Enhanced novel-view rendering without a 2D UIE model.

Overview of the 3D-USE pipeline

🛠️ Installation

git clone https://github.com/bilityniu/3D-USE.git
cd 3D-USE

conda env create -f environment/environment.yml
conda activate 3duse
bash environment/install.sh

📂 Data

Each scene contains COLMAP data and one Depth Anything V2 pseudo-depth map per image:

SCENE/
|-- images_wb/
|-- colmap/sparse/0/
`-- depth/

Generate pseudo-depth from the COLMAP input images:

python scripts/generate_da2_pseudodepth.py \
  --images /path/to/SCENE/images_wb \
  --output /path/to/SCENE/depth \
  --checkpoint /path/to/depth_anything_v2_vitl.pth \
  --da2-repo /path/to/Depth-Anything-V2

Image and depth filenames must have matching stems. Model weights are provided separately.

📦 Pretrained Models

Model Role Download
Depth Anything V2 Generates the monocular pseudo-depth used in Stage 1 Upstream release
UIE proposer Produces frozen view-wise proposals for ATC target construction Download
Transition calibrator Maps proposal transitions into the paired-data transition domain Download

The released proposer and calibrator can be used directly. See docs/TRAINING.md to prepare paired transition data and retrain the calibrator.

🚀 Training

bash recipes/run_two_stage.sh \
  --data /path/to/SCENE \
  --output /path/to/outputs/SCENE \
  --calibrator weights/transition_calibrator.pth \
  --uie-proposer weights/uie_proposer.pth

Direct interfaces: ns-train 3duse-stage1 --help and ns-train 3duse-stage2 --help.

🎥 Rendering and Evaluation

ns-render dataset \
  --load-config /path/to/stage2/config.yml \
  --output-path /path/to/renders/test \
  --split test \
  --rendered-output-names stage1_rgb enhanced_rgb
ns-eval \
  --load-config /path/to/stage1/config.yml \
  --output-path /path/to/reconstruction_metrics.json

Outputs include reconstruction, object, medium, depth, and enhanced branches.

Render a camera-path video:

ns-render camera-path \
  --load-config /path/to/stage2/config.yml \
  --camera-path-filename /path/to/camera_path.json \
  --output-path /path/to/enhanced.mp4 \
  --rendered-output-names enhanced_rgb

📖 Citation

@article{yuan2026three_duse,
  title={3D-USE: From Image-Level to Scene-Level Underwater Enhancement},
  author={Yuan, Jieyu and Zhang, Yuanlin and Li, Jihong and Guo, Chun-Le and Lu, Huimin and Li, Chongyi},
  journal={arXiv preprint arXiv:2608.28020},
  year={2026},
  eprint={2608.28020},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2608.28020}
}

🤝 Acknowledgements

3D-USE builds on Nerfstudio, gsplat, Plenodium, and WaterSplatting. The UIE proposer architecture follows UnderwaterRanker, and pseudo-depth is generated with Depth Anything V2. We thank the authors for making their work publicly available.

📜 License

Original 3D-USE contributions are released under the Pi-Lab License 1.0 for non-commercial research use. Third-party components and adapted files retain their original terms.

📮 Contact

Feel free to contact us at jieyuyuan.cn[AT]gmail.com for any questions or collaborations!

Releases

Packages

Contributors

Languages