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 Setup

Prerequisites

Python >= 3.8
PyTorch >= 1.7.0
torchvision >= 0.8.0
CUDA >= 10.2 (recommended)

Installation

# Clone the repository
git clone https://github.com/yourusername/SeaDiff.git
cd SeaDiff

# Create conda environment
conda create -n seadiff python=3.8
conda activate seadiff

📂 Dataset Preparation

To train SeaDiff, please follow these steps:

Download UIE datasets:
- UIEB
- EUVP
- SUIM-E
Generate depth maps using Depth Anything V2:

Create histogram representations:

python utils/create_hist_sample.py --input_dir datasets/UIEB/train/input --output_dir datasets/UIEB/train/histo

Dataset Structure

After preprocessing, organize your data as follows:

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

🚀 Quick Start

Training

Modify the configuration in conf.yml:

MODE: 1                    # 1 for training, 0 for inference
PRE_ORI: 'True'            # True for $x_0$, False for $\epsilon$
# ... other parameters

Start training:
```
python main.py
```

🏗️ Model Architecture

Overview of SeaDiff architecture

📜 Citation

If you find our work useful, please cite:

🤝 Acknowledgements

Our code is based on the following excellent works:

We thank the authors for their outstanding contributions! 🙏

📧 Contact

If you have any questions, please feel free to:

📧 Email: bihengyue@stu.ouc.edu.cn
🐛 Open an Issue
💬 Start a Discussion

⭐ If you find this project helpful, please consider giving it a star! ⭐