UDnet: underwater image enhancement

Pretrained weights for UDnet, from Adaptive deep learning framework for robust unsupervised underwater image enhancement (Saleh et al., Expert Systems with Applications, 2025).

• Code: https://github.com/alzayats/UDnet
• Paper: https://www.sciencedirect.com/science/article/pii/S0957417424031816
• Gallery: https://alzayats.github.io/UDnet/

UDnet corrects colour casts and improves contrast in underwater photos and keeps the input resolution.

Model	Status	File	Notes
UDnet V1	Released	UDnet.pth (5.6 MB)	The paper model. Fully unsupervised: trained without any reference images. 1.4 M parameters, runs on a CPU.
UDnet V2	Released	UDnet_v2.pth (26 MB)	The UDnet design (conditional VAE with latent feature modulation) modernised with NAFNet blocks and trained on paired images (UIEB and LSUI training splits) for stronger colour and haze correction. 12.8 M parameters (10.9 M at inference), best on a GPU.
V2 preview	Archived	UDnet_v2_preview.pth (59 MB)	The earlier V2 preview, a NAFNet U-Net (29 M parameters) trained the same way. Kept for reproducibility.

Usage

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 detail, 2x upscaled

The weights are downloaded from this repository automatically on first use. --ai adds a Real-ESRGAN pass that sharpens and upscales; it can invent fine texture, so use it for viewing, not measurement.

Files

File	Description
UDnet.pth	V1 inference weights (PyTorch state_dict, 1.4 M parameters)
UDnet_v2.pth	V2 weights ({'state_dict', 'config'}, float16, 12.8 M parameters including the training-only posterior encoder)
UDnet_v2_preview.pth	Archived V2 preview weights (float16, 29 M parameters); load_udnet_v2 loads either file

V2 comes within 0.2–0.5 dB PSNR of the 29 M-parameter preview on UIEB, LSUI, EUVP and UFO-120, ties on LPIPS and no-reference scores, and runs about 1.9× faster. Both are intended for research use.

Citation

@article{SALEH2025126314,
  title   = {Adaptive deep learning framework for robust unsupervised underwater image enhancement},
  journal = {Expert Systems with Applications},
  volume  = {268},
  pages   = {126314},
  year    = {2025},
  doi     = {https://doi.org/10.1016/j.eswa.2024.126314},
  author  = {Alzayat Saleh and Marcus Sheaves and Dean Jerry and Mostafa {Rahimi Azghadi}},
}