---
license: mit
library_name: pytorch
pipeline_tag: image-to-image
tags:
  - underwater-image-enhancement
  - image-enhancement
  - unsupervised-learning
---

# 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

```bash
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](https://github.com/xinntao/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 runs about 1.9× faster than the 29 M-parameter preview. Both are intended for research use.

## Citation

```bibtex
@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}},
}
```