UDBE: Unsupervised Diffusion-based Brightness Enhancement in Underwater Images

Tatiana Tais Schein, [email protected] | Gustavo Pereira de Almeira, [email protected] | Stephanie Loi Briao, [email protected] | Felipe Gomes de Oliveira, [email protected] | Rodrigo Andrade de Bem, [email protected] | Paulo L. J. Drews-Jr, [email protected]

Overview

Activities in underwater environments are paramount in several scenarios, which drives the continuous development of underwater image enhancement techniques. A major challenge in this domain is the depth at which images are captured, with increasing depth resulting in a darker environment. Most existing methods for underwater image enhancement focus on noise removal and color adjustment, with few works dedicated to brightness enhancement. This work introduces a novel unsupervised learning approach to underwater image enhancement using a diffusion model. Our method, called UDBE, is based on conditional diffusion to maintain the brightness etails of the unpaired input images. The input image is combined with a color map and a SignalNoise Relation map (SNR) to ensure stable training and prevent color distortion in the output images. The results demonstrate that our approach achieves an impressive accuracy rate in the datasets UIEB, SUIM and RUIE, well-established underwater image benchmarks. Additionally, the experiments validate the robustness of our approach, regarding he image quality metrics PSNR, SSIM, UIQM, and UISM, indicating the good performance of the brightness enhancement process.

Data Structure

UWSData
├── RUIE
│    ├── UCCS
│    |    ├── train
│    |    └── val
│    ├── UIQS
│    |    ├── train
│    |    └── val
│    └── UTTS
│         ├── train
│         └── val
├── SUIM
│    |── train
|    |    ├── images
|    |    └── masks
│    └── val
|         ├── images
|         └── masks
└── UIEB
     ├── train
     └── val

Setup

''' git clone https://github.com/gusanagy/UDBE.git
'''

''' cd UDBE '''

''' conda env create -f UDBE.yaml --name UDBE conda activate UDBE '''

Usage

Our code was trained on a computer with two NVIDIA TITAN X GPUs, totaling 24GB of GPU memory.

python #train from scratch, you can change setting in modelConfig 
python main.py --dataset_path "" --dataset "UIEB" -- state "train" 
python main.py --dataset_path "" --dataset "SUIM" -- state "train" 

python main.py --pretrained_path  --dataset "SUIM" -- state "val" --pretrained_path "1000.pt"
python main.py --pretrained_path  --dataset "UIEB" -- state "val" --pretrained_path "1000.pt"

python main.py --pretrained_path  --dataset "SUIM" -- state "inference" --inference_image " "  --pretrained_path "1000.pt"
python main.py --pretrained_path  --dataset "UIEB" -- state "inference" --inference_image " " --pretrained_path "1000.pt"

Tests can be done in the notebook avaliacao.ipynb. As well as viewing images from the respective datasets. When running this notebook there are options to download and automatically generate the folders, datasets and checkpoints necessary to run the model.

Paper: https://arxiv.org/abs/2501.16211 For quotes:

''' @article{schein2025udbe, title={UDBE: Unsupervised Diffusion-based Brightness Enhancement in Underwater Images}, author={Schein, Tatiana Ta{'\i}s and de Almeira, Gustavo Pereira and Bri{~a}o, Stephanie Loi and de Bem, Rodrigo Andrade and de Oliveira, Felipe Gomes and Drews-Jr, Paulo LJ}, journal={arXiv preprint arXiv:2501.16211}, year={2025} } '''

ACKNOWLEDGMENT

This study was funded, in part, by the São Paulo Research Foundation (FAPESP), Brazil, under Process Number 2024/10523-5. The authors would also like to thank the PRH-ANP and CNPQ organizations for their research support and financial assistance.