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