2021
U-shape Transformer for Underwater Image Enhancement
Lintao Peng, Chunli Zhu, Liheng Bian*
1Beijing Institute of Technology
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 constructed a large-scale underwater image (LSUI) dataset including 4279 image pairs, and 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, 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.
[arXiv] [Code] [Data]
Highlights
1. We reported a novel U-shape Transformer dealing with the UIE task, in which the designed channel-wise and spatial-wise attention mechanism enables to effectively remove color artifacts and casts.
2. We designed a novel multi-color space loss function combing the RGB, LCH and LAB color-space features, which further improves the contrast and saturation of output images.
3. We released a large-scale underwater dataset containing 4279 image pairs, which facilitates further development of underwater imaging techniques.
4. Extensive experiments have show that the U-shape transformer we proposed combined with the multi-color space loss function achieves state-of-the-art performance on several public datasets and our dataset.
We released a large-scale dataset(LSUI) containing 4279 real underwater images and the corresponding high-quality reference images, semantic segmentation maps, and medium transmission maps, which involve richer underwater scenes (lighting conditions, water types and target categories) and facilitates further development of UIE techniques. You can download it from [BaiduYun](password is lsui) or [GoogleDrive]. If you want to use the LSUI dataset, please cite our [paper]
Overall Architecture
Fig 1. Detailed network structure. The model is modified based on Unet-Gan. In order to guide the network to pay attention to the more severely attenuated color channels and spatial regions, we introduced the CMSFFT module and SGFMT module based on transformer architecture . In order to make the training process of our nodel more stable, we added a multi-scale gradient flow mechanism to the network.
Detailed structure of SGFMT and CMSFFT
Fig 2. SGFMT data flow diagram.
Fig 3. CMSFFT detailed structure
Results
Fig 4. Enhancement results of U-shape transformer trained on different underwater datasets. (a): Underwater images sampled from Test-L504; (b): Enhanced results using the model trained on the underwater dataset Train-U2050; (c): Enhanced results using the model trained on the underwater dataset Train-E11435; (d): Enhanced results using the model trained by our proposed dataset Train-L4500; (e): References images (Regarded as Ground Truth)
Fig 5. Visual comparison of enhancement results sampled from Test-L504. From left to right are raw images, and the results of UIBLA, Retinex based, FUnLE, UGAN, UIE-DAL, Ucolor, and our U-shape Transformer.
Fig 6. Visual comparison of enhancement results sampled from Test-U60. From left to right are raw images, and the results of UIBLA , Retinex based, FUnLE, UGAN, UIE-DAL, Ucolor, and our U-shape Transformer.
Tab 7. Evaluation with Reference Images on Test-L504 and Test-U90
Tab 8. Evaluation without Reference Images on Test-U60 and SQUID
Technical Paper
Demo Video
Citation
Lintao Peng, Chunli Zhu, and Liheng Bian. "U-shape Transformer for Underwater Image Enhancement", 2021
@misc{peng2021ushape,
title={U-shape Transformer for Underwater Image Enhancement},
author={Lintao Peng and Chunli Zhu and Liheng Bian},
year={2021},
eprint={2111.11843},
archivePrefix={arXiv},
primaryClass={cs.CV}
}