SeaDiff: Underwater Image Enhancement with Degradation-Aware Diffusion Model

This is the official repository for the paper [SeaDiff: Underwater Image Enhancement with Degradation-Aware Diffusion Model].

News

• 2025.06.12: The initial version of the code is uploaded.

Environment

• python >= 3.8
• pytorch >= 1.7.0
• torchvision >= 0.8.0

Dataset Preparation

To train SeaDiff, you should:

1. Download the UIE datasets.

2. Then use Depth Anything to estimate monocular depth maps.

3. Third, use utils/create_hist_sample.py to estimate histogram representations.

After preprocessing, our folder structure is as follows:

datasets/
└── UIEB/
    ├── train/
    │   ├── input/
    │   ├── label/
    │   ├── depth/
    │   └── histo/
    └── val/
        ├── input/
        ├── label/
        ├── depth/
        └── histo/

🌟 Training and 🎇 Testing

Whether it's for training or inference, you just need to modify the configuration parameters in conf.yml and run main.py. MODE=1 is for training, MODE=0 is for inference.

📜 Citation

If you find our work useful, please cite:

🤝 Acknowledgements

Our code is based on DocDiff, HistoGAN and Depth Anything. We thank the authors for their excellent work!

If you have any questions, please don't hesitate to open an issue or contact Hengyue Bi at bihengyue@stu.ouc.edu.cn. 🤞🤞🤞