Underwater image enhancement · open source
Bring the colour back to underwater photos.
Water soaks up red light within a few metres, so most underwater footage comes out green or blue and hazy. UDnet is a neural network that restores natural colour and contrast in one pass. Drag the handle to compare.
The problem
Red goes first.
Each colour fades at a different depth. Red is mostly gone by about 5 m, orange by 10 m and yellow by 20 m, while green and blue carry much further. Floating particles add a milky haze on top. UDnet has seen thousands of these scenes and learned how to put the missing colour back.
Approximate, for clear ocean water. Coastal and green water lose light much faster.
Gallery
24 scenes from public test sets
None of these images were used to train V2. Pick a version once and every slider follows.
Versions
Pick the version that fits the job
UDnet V1
The model from our ESWA 2025 paper. Fully unsupervised: it learned from raw underwater photos alone, with no reference images, so it never copies another method's style.
- 5.6 MB of weights, runs on a laptop CPU
- Gentle, conservative correction
- Weights auto-download on first run
UDnet V2
Trained on thousands of paired images with a modern restoration network. Much stronger colour and haze correction while staying faithful to the scene.
- Best accuracy against reference images
- Needs a GPU for large images
- Same command, add
--model v2
V2 + AI detail
Runs V2, then a generative model that sharpens and upscales 2×. Great for viewing and sharing.
- Crisp edges, less noise and blocking
- Can invent fine texture that was not captured
- Not for measurements or science
V2 + AI Max
Runs V2, then a diffusion model (DiffBIR) that restores fine natural detail such as spines, scales and tentacles, upscaled 2×. Colours stay V2's.
- The most detailed look
- About 25 s per image on a GPU
- Invents more than V2 + AI: viewing only
- Set up from the ai-max branch
Measured
Scores on held-out test images
PSNR (dB) and SSIM measure how close an output is to each test set's reference image; higher is better. V1 is fully unsupervised and never saw a reference image, so it is not trying to match them. V2 was trained on paired images, which is what these scores reward. V2 figures are from the current preview model.
| Test set | Images | V1 PSNR | V2 PSNR | V1 SSIM | V2 SSIM |
|---|---|---|---|---|---|
| UIEB test split | 90 | 19.86 | 22.86 | 0.842 | 0.909 |
| LSUI test split | 400 | 19.64 | 28.07 | 0.807 | 0.910 |
| EUVP test samples | 515 | 19.59 | 25.93 | 0.783 | 0.874 |
| UFO-120 test | 120 | 18.53 | 24.71 | 0.628 | 0.734 |
Use it
Enhance a folder in one command
Works on a single image or a whole folder. The weights download automatically the first time.
git clone https://github.com/alzayats/UDnet.git cd UDnet && pip install -r requirements.txt python Inference.py --input path/to/images --output results/ # V1 python Inference.py --input path/to/images --output results/ --model v2 # V2 python Inference.py --input path/to/images --output results/ --model v2 --ai x2 # V2 + AI
Cite UDnet
@article{SALEH2025126314,
title = {Adaptive deep learning framework for robust unsupervised underwater image enhancement},
author = {Alzayat Saleh and Marcus Sheaves and Dean Jerry and Mostafa {Rahimi Azghadi}},
journal = {Expert Systems with Applications},
volume = {268},
pages = {126314},
year = {2025},
doi = {10.1016/j.eswa.2024.126314}
}