U-Shape Transformer for Underwater Image Enhancement

TLDR

Underwater images suffer from light absorption, scattering, and uneven color attenuation, and existing datasets lack diversity and high‑fidelity references. The authors create a large‑scale underwater image dataset and propose a U‑shape Transformer network for image enhancement. The dataset contains 4,279 paired raw, reference, segmentation, and transmission images, while the U‑shape Transformer incorporates CMSFFT and SGFMT modules and a multi‑color‑space loss to focus on color and spatial attenuation. Experiments show the method outperforms state‑of‑the‑art by over 2 dB, and the dataset and code are publicly released.

Abstract

The light absorption and scattering of underwater impurities lead to poor underwater imaging quality. The existing data-driven based underwater image enhancement (UIE) techniques suffer from the lack of a large-scale dataset containing various underwater scenes and high-fidelity reference images. Besides, the inconsistent attenuation in different color channels and space areas is not fully considered for boosted enhancement. In this work, we built a large scale underwater image (LSUI) dataset, which covers more abundant underwater scenes and better visual quality reference images than existing underwater datasets. The dataset contains 4279 real-world underwater image groups, in which each raw image's clear reference images, semantic segmentation map and medium transmission map are paired correspondingly. We also reported an U-shape Transformer network where the transformer model is for the first time introduced to the UIE task. The U-shape Transformer is integrated with a channel-wise multi-scale feature fusion transformer (CMSFFT) module and a spatial-wise global feature modeling transformer (SGFMT) module specially designed for UIE task, which reinforce the network's attention to the color channels and space areas with more serious attenuation. Meanwhile, in order to further improve the contrast and saturation, a novel loss function combining RGB, LAB and LCH color spaces is designed following the human vision principle. The extensive experiments on available datasets validate the state-of-the-art performance of the reported technique with more than 2dB superiority. The dataset and demo code are available at https://bianlab.github.io/.