This repository contains the official implementation of the paper Style-Decoupled Adaptive Routing Network for Underwater Image Enhancement arXiv
- python == 3.9
- PyTorch == 2.4.0
- CUDA == 11.8
- ninja == 1.11.1
- einops
- diffusers == 0.36.0
- tqdm
- numpy
- matplotlib
The code has been tested on Ubuntu 20.04 with NVIDIA GeForce RTX 3080 Ti, 4090 and 5090 D. Its dependences version is not very strict, but I recommend at least keeping CUDA==11.8.
We provide the divided dataset for UIEB and LSUI. This division refers to WF-diff
| UIEB | UIEB |
| LSUI | LSUI |
Extract code:123x
Different dataset splits can cause fluctuations in training and testing results.
# Specify the dataset path and checkpoint path in test.py
python test.py
The network module contains layers with random operations (NonLocalSparseAttention), so the test result metrics may have slight fluctuations.
# Specify the dataset path and checkpoint path in train.py
python train.py
The repository is based on U-Shape, some of the code is borrowed from:
Thanks for their opensourceing.