Deep dive into clarity: Leveraging signal-to-noise ratio awareness and knowledge distillation for underwater image enhancement

Guodong Fan, Jingchun Zhou, Chengpei Xu, Zheng Cheng

Abstract: This paper presents an innovative dual-branch solution designed for underwater image enhancement (UIE), leveraging the synergistic combination of Signal-to-Noise Ratio (SNR) aware transformers and convolutional models. SNR-Net dynamically enhances pixel quality through spatial-varying operations. While transformers excel in capturing long-range dependencies, they face challenges in weak local relation learning. To address this, we introduce a SNR prior to guide transformer learning, incorporating a novel self-attention mechanism that avoids tokens from regions with very low SNR. Conversely, CNNs, optimized for exploiting local patterns, suffer from limited receptive fields and weak diversity representation. To overcome this limitation, we enhance the receptive field and multi-scale perception of CNNs by introducing a MR-ResNet module. Additionally, we incorporate a Selective Kernel Merging Module (SKMM), an attention-based feature merging module. These enhancements empower our approach to learn an enriched set of features that selectively combine contextual information from both branches while preserving high-quality spatial details. Finally, through knowledge distillation and contrastive learning, SNR-KD significantly reduces the number of parameters and computations of SNR-Net with minimal impact on performance. Extensive experiments validate the effectiveness of our methods, namely SNR-Net and SNR-KD, demonstrating their state-of-the-art performance compared to other recent UIE methods. The code of our model is publicly available at: https://github.com/Alexande-rChan/SNR-UIE.