SwiftWater: A lightweight SwiftFormer-based underwater image enhancement network with auxiliary prompt module

Abstract

Due to variations in light attenuation underwater and the complexities of light propagation in aquatic environments, underwater images commonly suffer from degradation, such as color distortion, blur, and low contrast. These challenges pose significant obstacles to researchers studying marine resources. In recent years, numerous solutions have been proposed for underwater image enhancement (UIE) tasks, with deep learning-based methods, including those employing convolution and transformer structures, demonstrating promising performance. However, deep learning approaches often encounter issues related to large parameter sizes and substantial computational demands. Addressing these concerns, this paper introduces SwiftWater, a lightweight transformer-based network specifically designed for UIE tasks, featuring a U-Net style architecture. Leveraging the efficient SwiftFormer encoder as the network backbone, we introduce a Hybrid Squeeze and Excitation Block (HSEB) to dynamically adjust representations within skip connections. Additionally, we introduce an Auxiliary Prompt Module (APM) with a Lightweight Prompt Block (LPB) that utilizes the HSV color space of underwater images to enhance training. LPB within APM comprises two modules: the Prompt Generation Module (PGM) for prompt generation and the Lightweight Prompt Interaction Module (LPIM) for prompt-feature interaction, generating rich auxiliary image representations. Compared to existing mainstream UIE methods employing transformer structures, our proposed approach offers reduced parameter counts while achieving state-of-the-art (SOTA) performance across both quantitative and qualitative metrics.