Image-Conditional Diffusion Transformer for Underwater Image Enhancement

TL;DR

This paper introduces a novel underwater image enhancement method using an image-conditional diffusion transformer that outperforms existing techniques in quality and scalability.

Contribution

It proposes replacing the U-Net backbone with a transformer in a diffusion model for improved underwater image enhancement.

Findings

ICDT-XL/2 achieves state-of-the-art enhancement quality.

The hybrid loss accelerates the sampling process.

The method demonstrates excellent scalability on the Underwater ImageNet dataset.

Abstract

Underwater image enhancement (UIE) has attracted much attention owing to its importance for underwater operation and marine engineering. Motivated by the recent advance in generative models, we propose a novel UIE method based on image-conditional diffusion transformer (ICDT). Our method takes the degraded underwater image as the conditional input and converts it into latent space where ICDT is applied. ICDT replaces the conventional U-Net backbone in a denoising diffusion probabilistic model (DDPM) with a transformer, and thus inherits favorable properties such as scalability from transformers. Furthermore, we train ICDT with a hybrid loss function involving variances to achieve better log-likelihoods, which meanwhile significantly accelerates the sampling process. We experimentally assess the scalability of ICDTs and compare with prior works in UIE on the Underwater ImageNet dataset. Besides good scaling properties, our largest model, ICDT-XL/2, outperforms all comparison methods, achieving state-of-the-art (SOTA) quality of image enhancement.