Source-linked AI summary
Shallow-UWnet : Compressed Model for Underwater Image Enhancement
Ankita Naik, Apurva Swarnakar, Kartik Mittal
TL;DR
Underwater image enhancement models can be computationally and memory intensive, limiting deployment on portable devices. Shallow-UWnet uses a shallow architecture to maintain comparable performance with fewer parameters, while generalizing across varied synthetic and real-world datasets.
Problem
Deep CNN- and GAN-based underwater enhancement models require substantial computation and memory, creating a deployment challenge for real-time applications.
Method
Shallow-UWnet is a shallow convolutional architecture composed of densely connected convolutional blocks with skip connections.
Results
Shallow-UWnet maintains comparable quantitative performance while using 18 times fewer trainable parameters and testing 10 times faster.
Takeaways & Limitations
The model generalizes across varied datasets and supports lighter, faster on-device and real-time underwater image enhancement.
Abstract
from arXiv · showhide
Over the past few decades, underwater image enhancement has attracted increasing amount of research effort due to its significance in underwater robotics and ocean engineering. Research has evolved from implementing physics-based solutions to using very deep CNNs and GANs. However, these state-of-art algorithms are computationally expensive and memory intensive. This hinders their deployment on portable devices for underwater exploration tasks. These models are trained on either synthetic or limited real world datasets making them less practical in real-world scenarios. In this paper we propose a shallow neural network architecture, \textbf{Shallow-UWnet} which maintains performance and has fewer parameters than the state-of-art models. We also demonstrated the generalization of our model by benchmarking its performance on combination of synthetic and real-world datasets.
Introduction
Underwater robotics increasingly relies on real-time image interpretation, but wavelength-dependent absorption and scattering degrade underwater imagery and hinder downstream vision tasks. Image enhancement is therefore presented as a necessary preprocessing step, while existing deep CNN- and GAN-based methods face computational and memory limitations.
- Motivation: Underwater robotics applications increasingly use autonomous underwater vehicles for surveillance, mapping, exploration, rescue, collection, and military operations.Many applications require real-time interpretation of images or videos for environmental perception.
- Image degradation: Wavelength-dependent absorption and particle scattering reduce underwater image visibility and contrast while causing color deviations and green or blue casts.These degradations limit image applicability for downstream tracking, classification, and detection.
- Image enhancement: Image improvement precedes downstream underwater interpretation, and enhancement restores perceptual and statistical image qualities in real time.The study focuses on enhancement rather than restoration or supplementary task-specific methods.
- Existing methods: Existing very deep CNN- and GAN-based enhancement models address noise, contrast, multi-information improvement, and dehazing but require high computational and memory resources.Enhancement methods are described as more generalized than restoration methods because they extract image information without prior environmental knowledge.
Background
Underwater image-enhancement methods comprise non-physical, physical model-based, and deep-learning approaches, but the first two overlook specific underwater properties. Deep-learning methods focus on color correction, while reviewed GAN- and CNN-based models retain important limitations.
- Background: Image-enhancement methods are broadly classified as non-physical, physical model-based, or deep-learning methods.
- Background: Non-physical and physical models are insufficient for underwater enhancement because they ignore specific underwater properties.Non-physical models adjust pixel values, whereas physical models estimate parameters describing image degradation.
- Background: Deep-learning methods tend to perform better by focusing on color correction and primarily using GANs and CNNs.The paper reviews current state-of-the-art GAN- and CNN-based models in detail.
- Background: FUnIE-GAN uses conditional adversarial training for image-to-image translation but mismodels sunlight, amplifies background noise, and can produce over- or under-saturated images.It assumes a nonlinear mapping between distorted and enhanced images and trains on a large-scale dataset.
- Background: Water-Net is a gated-fusion CNN using white balance, gamma correction, and histogram equalization inputs, but its complex architecture suffers from backscatter.It is trained on UIEB for underwater image enhancement.
- Background: UResnet is a CNN-based residual network providing a more comprehensive supervised-learning method for underwater image enhancement.
Underwater Datasets
The study benchmarks Shallow-UWnet’s generalization across three underwater image datasets spanning synthetically distorted and real-world imagery. These datasets differ in pairing, capture conditions, and evaluation scope.
- Dataset selection: Three datasets are used to benchmark model generalization against state-of-the-art models: EUVP, UIEB, and UFO-120.They span synthetically distorted and real-world underwater images.
- EUVP Dataset: EUVP contains 10K paired and 25K unpaired images of poor and good perceptual quality captured during oceanic explorations under varied visibility conditions.Its paired images were generated by distorting real-world images with a CycleGAN-based underwater distortion model.
- UIEB Dataset: UIEB comprises 890 real underwater images captured under different lighting conditions, with diverse color ranges and contrast levels.Reference images are color-cast-free and improve visibility and brightness relative to source images.
- UFO-120 Dataset: UFO-120 comprises 1620 paired underwater images collected during explorations in different water types, with distorted counterparts generated using style-transfer techniques.The study evaluates on the 120 paired underwater images labeled as the testing set.
Evaluation Metrics
The evaluation uses reference-based PSNR and SSIM, non-reference UIQM, and compression and speedup rates to assess enhancement quality, model compression, and acceleration. UIQM is defined as a weighted combination of UICM, UISM, and UIConM, with coefficients specified from prior work.
- Image-quality metrics: PSNR and SSIM quantify reconstruction quality and structural similarity against the respective reference image.These are standard quantitative metrics used to evaluate the model’s output images.
- Image-quality metrics: UIQM provides a non-reference measure of underwater image quality.The passage identifies UIQM as an additional quality analysis measure for generated output images.
- Image-quality metrics: UIQM = c1 × UICM + c2 × UISM + c3 × UIConM.The coefficients are set to c1 = 0.0282, c2 = 0.2953, and c3 = 3.5753 according to Panetta, Gao, and Agaian (2015).
- Compression and acceleration: Compression and speedup rates measure the quality of model compression and acceleration.The compression notation uses α(M) for model parameters and β(M) for testing time per image, comparing original model M with compressed model M∗.
Our Proposed Approach
Shallow-UWnet uses a shallow architecture with three densely connected convolutional blocks and skip connections from the raw input image. It is trained with combined pixel-wise MSE and VGG perceptual losses to enhance underwater images while preserving structural and texture properties.
- Architecture: Shallow-UWnet processes 256 × 256 RGB underwater images through an initial convolution layer and three chained convolution blocks, followed by a final layer producing the enhanced image.The initial 3x3 convolution generates 64 feature maps with ReLU activation, while the final convolution uses 3 kernels.
- Architecture: The architecture concatenates the raw input image with each block’s output through skip connections, supporting feature learning and preserving essential raw-image characteristics.These connections are also described as helping address vanishing gradients by weighting raw-input channels more strongly.
- Architecture: Each ConvBlock contains convolution, dropout, and ReLU layers, while the block sequence and skip connections deter overfitting and support network generalization.A further Conv-ReLU pair facilitates concatenation with the raw image from the skip connection.
- Loss Function: The model uses a multi-term loss combining pixel-wise MSE with VGG perceptual loss to preserve edge sharpness and impose structural and texture similarity.MSE compares the estimated image with the clear ground truth, while VGG loss compares feature representations from the last convolutional layer of a pretrained 19-layer VGG network.
- Loss Function: The final loss L is the summation of the MSE and VGG perceptual losses.The perceptual loss is based on ReLU activations from the pretrained VGG network.
Experimental Evaluations
The model was evaluated through qualitative and quantitative comparisons with recent state-of-the-art underwater image enhancement methods on synthetic and real-world images.
- Experimental Evaluations: Evaluations compared Shallow-UWnet qualitatively and quantitatively against WaterNet, FunIE-GAN, and Deep SESR on synthetic and real-world underwater images.The comparisons used checkpoints provided by the corresponding authors to produce the best results for objective evaluation.
1. Training and Validation Data
Shallow-UWnet was trained and validated on EUVP Underwater ImageNet, chosen for its diverse camera sources, visibility conditions, locations, and perceptual quality to support generalization.
- 1. Training and Validation Data: EUVP Underwater ImageNet provided paired training and validation images collected with varied cameras under diverse visibility conditions.The paired images were generated using CycleGAN, and the dataset was selected for its diversity across locations and perceptual quality.
2. Network Implementation and Training
Shallow-UWnet was trained with specified optimization settings in PyTorch and evaluated on diverse synthetic and real-world underwater datasets to assess transferability.
- Network Implementation and Training: ADAM training used a 0.0002 learning rate, 0.2 layer dropout, and batch size 1.The model was optimized for 50 epochs, taking around ten hours per model.
- Network Implementation and Training: 50 epochs required around ten hours on an Intel i7-8750H CPU, 16GB RAM, and Nvidia GTX 1060 GPU using PyTorch.The stated training environment included 16GB RAM and a Nvidia GTX 1060 GPU.
- Testing Data Sets: The model was tested on synthetic and real-world images from varied datasets to benchmark transferability.The testing datasets included UFO-120, EUVP Dark, and UIEB.
- Testing Data Sets: 120 UFO-120 images generated from clear ocean-exploration images across different water types formed one testing subset.Style transfer generated the corresponding distorted images.
- Testing Data Sets: 1000 EUVP Dark paired images with dark underwater backgrounds were used for testing from a dataset clustered with 5500 paired images.The paper selected a subset of 1000 images.
- Testing Data Sets: UIEBD supplied 890 paired underwater images captured under different light conditions and spanning diverse color ranges and contrast levels.Reference images were generated through meticulous pairwise comparisons.
Results
Shallow-UWnet achieves comparable performance across three test datasets while using fewer trainable parameters and processing images faster. On EUVP-Dark, it outperforms other models on PSNR, SSIM, and UIQM, indicating generalization beyond its training conditions.
- Overall results: Shallow-UWnet achieves comparable performance across all three test datasets while using fewer trainable parameters and processing test images faster.These properties support lighter on-device deployment and real-time underwater image enhancement applications.
- EUVP-Dark: On EUVP-Dark, Shallow-UWnet outperforms all other models on PSNR, SSIM, and UIQM.It cleans color hue and sharpens images despite training on images with better lighting conditions, emphasizing its generalizing nature.
- UFO-120: On UFO-120, Deep SESR has the best performance of all models because it was trained on UFO-120.
Conclusion
Shallow-UWnet maintains comparable quantitative performance while using fewer trainable parameters, testing faster, and generalizing across varied datasets for real-world underwater imaging.
- Shallow-UWnet maintains comparable quantitative performance for underwater image enhancement.
- 18 times lesser trainable parameters and 10 times faster testing improve the model’s deployment efficiency.
- Shallow-UWnet generalizes on varied datasets, emphasizing its real-world application.