UIEFormer: Lightweight Vision Transformer for Underwater Image Enhancement

Abstract

The selective absorption and scattering of light in water degrade underwater image quality, hindering the performance of underwater tasks. Moreover, existing data-driven underwater image enhancement (UIE) methods rely on large-scale, high-quality underwater image data sets, which are costly to acquire in terms of time and labor. In this work, we present a UIE framework named UIEFormer, which is built upon a popular conventional image defogging framework DehazeFormer, possessing satisfactory performance on a small-scale training data set of underwater images. We propose an interpolation-based upsampling strategy to avoid checkerboard artifacts caused by PixelShuffle. Extra feature channels are introduced to segregate noncritical high-level image features for UIE tasks. Further, we apply a loss function combining per-pixel loss, perceptual loss, and coloration loss to adapt to the underwater environment. Results on real-world data sets demonstrate that our method has certain advantages over classical and popular UIE methods. In addition, we conduct ablation experiments to demonstrate the contribution of each module in our work. We also demonstrate the practical significance of our approach for underwater image processing tasks.