AquaAE: A Lightweight Deep Learning Network for Underwater Image Restoration

Due to the physical properties of water and the presence of suspended particles, underwater images often exhibit a bluish-green tint, reduced contrast, and uneven light distribution. Many researchers strive for better image restoration techniques, but they often overlook the high computational demands of these models, limiting their application in resource-constrained scenarios. To address this, we have introduced a model named AquaAE for image restoration. This model adopts a simple autoencoder structure, utilizing skip connections to merge features from the encoder and decoder, and incorporates a red channel enhancement to improve image restoration quality. Compared to advanced deep learning networks like U-Transformer, Twin-UIE, and Semi-UIR, our model is more straightforward, employing only simple convolution and upsampling. Combined with our specially calculated red channel enhancement coefficients tailored for different water conditions, AquaAE efficiently captures local features and spatial relationships, thereby better restoring image colors.AquaAE excels in classical evaluation metrics such as Mean Squared Error (MSE), Peak Signal-to-Noise Ratio (PSNR), and Structural Similarity Index Measure (SSIM). More critically, our model demonstrates outstanding computational efficiency, with its FLOPs being only 2.6% of Twin-UIE's (1.32G) and its parameter count merely 2.8% of U-Transformer's (0.88M), highlighting the lightweight nature of the model. This lightweight design is crucial not only for improving image restoration effectiveness but also for underwater mobile devices with limited computational resources and battery life. We trained AquaAE on the underwater scenes subset of the EUVP dataset and tested it on the underwater ImageNet subset of the EUVP dataset. The results show that AquaAE performs exceptionally well in underwater image restoration.