Source-linked AI summary
Contrastive Semi-supervised Learning for Underwater Image Restoration via Reliable Bank
Shirui Huang, Keyan Wang, Huan Liu, Jun Chen, Yunsong Li
TL;DR
Underwater image restoration lacks abundant, well-matched labeled data, while mean-teacher training can suffer from incorrect pseudo-labels and confirmation bias. Semi-UIR addresses these issues with a reliable bank and contrastive regularization, and reports improvements over state-of-the-art methods on full-reference and non-reference benchmarks.
Problem
Paired real-world underwater images are difficult to acquire, and existing datasets suffer from synthetic-to-real discrepancy and limited size.
Method
Semi-UIR uses mean-teacher training with a monotonicity-selected NR-IQA reliable bank for pseudo-labels and contrastive regularization against confirmation bias.
Results
Semi-UIR reports obvious quantitative and qualitative improvements over state-of-the-art methods on full-reference and non-reference underwater benchmarks.
Takeaways & Limitations
Reliable teacher-student consistency and contrastive regularization are reported as the sources of Semi-UIR’s superior performance over other state-of-the-art algorithms.
Takeaways & Limitations
The teacher is not guaranteed to outperform the student consistently, so wrong pseudo-labels can jeopardize student training.
Abstract
from arXiv · showhide
Despite the remarkable achievement of recent underwater image restoration techniques, the lack of labeled data has become a major hurdle for further progress. In this work, we propose a mean-teacher based Semi-supervised Underwater Image Restoration (Semi-UIR) framework to incorporate the unlabeled data into network training. However, the naive mean-teacher method suffers from two main problems: (1) The consistency loss used in training might become ineffective when the teacher's prediction is wrong. (2) Using L1 distance may cause the network to overfit wrong labels, resulting in confirmation bias. To address the above problems, we first introduce a reliable bank to store the "best-ever" outputs as pseudo ground truth. To assess the quality of outputs, we conduct an empirical analysis based on the monotonicity property to select the most trustworthy NR-IQA method. Besides, in view of the confirmation bias problem, we incorporate contrastive regularization to prevent the overfitting on wrong labels. Experimental results on both full-reference and non-reference underwater benchmarks demonstrate that our algorithm has obvious improvement over SOTA methods quantitatively and qualitatively. Code has been released at https://github.com/Huang-ShiRui/Semi-UIR.
1. Introduction
Underwater restoration is constrained by severe visual degradation and scarce, mismatched paired data. Semi-UIR addresses these limitations by using unlabeled images with reliable teacher pseudo-labels and contrastive regularization.
- Motivation: Underwater images often exhibit color distortion, low contrast, and blur from refraction, absorption, and scattering.These defects can reduce visual appeal and potentially hinder underwater robotic systems.
- Motivation: Synthetic benchmarks differ from real underwater data in ambient light and color cast, creating domain shifts that weaken real-scene generalization.UWCNN uses an image formation model, whereas EUVP uses GAN synthesis; both differ from UIEB real-world images.
- Motivation: UIEB provides only 890 paired underwater images, leaving supervised models vulnerable to overfitting.Existing paired datasets are also limited because real paired underwater images are difficult to acquire.
- Approach: Semi-UIR adopts mean-teacher training to exploit relatively easy-to-collect unlabeled underwater images and improve real-world generalization.The teacher is formed from the student through exponential moving average, and its predictions guide student training.
- Approach: A reliable bank selects best-ever teacher outputs using an NR-IQA metric chosen through monotonicity analysis as pseudo-labels.This design targets unreliable teacher predictions and improves pseudo-label reliability.
- Approach: Contrastive loss is added as regularization to alleviate confirmation bias from incorrect pseudo-labels.The framework combines reliable pseudo-labeling with contrastive regularization as its principal contributions.
2. Related Work
Prior underwater restoration methods use handcrafted priors, physical models, or supervised mappings, while semi-supervised and contrastive approaches remain limited in this domain. The paper positions its method at the intersection of these research directions.
- Underwater Image Restoration: Traditional methods use handcrafted priors or enhancement techniques, including CLAHE, Retinex, fusion, and MMLE.Model-based methods estimate parameters such as transmission and ambient light, while model-agnostic methods design enhancement operations.
- Underwater Image Restoration: Early deep learning methods estimate physical imaging parameters, but inaccurate transmission and ambient-light estimates can hinder restoration quality.More recent methods directly learn restoration mappings from labeled datasets in a supervised manner.
- Semi-supervised Learning: Semi-supervised learning combines labeled and unlabeled data through methods such as mean teacher, virtual adversarial learning, and FixMatch.Mean teacher has been applied successfully to recognition, segmentation, and restoration, but remains rarely explored for underwater restoration.
- Contrastive Learning: Contrastive learning separates dissimilar samples and brings similar samples closer, but prior restoration uses generally depend on supervised clean and degraded examples.Applying contrastive loss to unlabeled underwater restoration remains an open problem identified by the paper.
3. Method
Semi-UIR uses a mean-teacher framework to learn from labeled and unlabeled underwater images. It improves pseudo-label reliability with a bank of best-ever teacher outputs and reduces confirmation bias through contrastive regularization.
- 3.2. Semi-supervised Underwater Restoration: Semi-UIR trains teacher and student networks with supervised and unsupervised objectives on disjoint labeled and unlabeled datasets.The student is updated by gradient descent, while the teacher provides supervision for unlabeled images.
- 3.2. Semi-supervised Underwater Restoration: The teacher is updated by exponential moving average, aggregating previously learned student weights to stabilize training.The update uses momentum η in the range (0, 1).
- 3.3. Reliable Teacher-Student Consistency: Wrong teacher predictions can jeopardize student training, so Semi-UIR replaces naive pseudo labels with best-ever outputs stored in a reliable bank.At each iteration, current teacher outputs are compared with student outputs and banked pseudo labels before replacement.
- 3.4. Reliable Metric Selection: MUSIQ is selected to assess teacher-output reliability because its scores best follow the expected monotonic quality trend among evaluated NR-IQA methods.The analysis uses fused images whose visual quality deteriorates as α increases.
- 3.5. Contrastive Regularization: Contrastive regularization addresses confirmation bias by treating reliable-bank outputs as positive samples and strongly augmented degraded images as negative samples.Feature-space L1 distances are computed using hidden representations from a pre-trained VGG-19.
4. Experimental Results
Semi-UIR is evaluated on full-reference and non-reference underwater benchmarks, with ablations examining its reliable bank, contrastive regularization, metric selection, and training behavior.
- Comparison with the State-of-the-Arts: The evaluation uses full-reference benchmarks with PSNR and SSIM, and non-reference benchmarks with UIQM, UCIQE, and MUSIQ.The compared methods are retrained on the authors’ training set.
- Comparison with the State-of-the-Arts: Qualitative comparisons show robust restoration across underwater scenes, with natural color and rich details.The authors attribute this behavior to reliable pseudo labels and contrastive regularization.
- Ablation Study: The full Semi-UIR solution performs best in the reported ablation results, while reliable-bank and contrastive-loss variants isolate their respective contributions.The ablation compares supervised training, mean-teacher training, reliable-bank use, contrastive loss, and the complete method.
- Ablation Study: Without reliable positive samples, contrastive loss causes over-enhancement; without contrastive loss, reliable-bank outputs retain color distortion and remain close to degraded inputs.Together, the qualitative ablations support using both components.
- Breakdown of the Training: At 10 epochs, the teacher prediction is much better than the student’s, and both outputs improve simultaneously as training proceeds.These intermediate results support using teacher outputs to train the student network.
- Influence of Non-reference Metric: MUSIQ provides the best performance among the evaluated NR-IQA choices for selecting reliable-bank pseudo labels.The comparison evaluates NIMA, PAQ2PIQ, and MUSIQ on testS and testR.
5. Conclusion
The paper concludes that Semi-UIR is an efficient semi-supervised underwater image restoration method whose performance is attributed to reliable teacher-student consistency and contrastive regularization.
- Semi-UIR’s superior performance over other state-of-the-art algorithms is attributed to reliable teacher-student consistency and contrastive regularization.Future work is proposed for other restoration tasks and improved memory management.
A. Supplementary Material
The supplementary material provides additional architectural, training, validation, metric, and benchmark analyses.
- The supplement details AIM-Net structure, validation data, loss functions, optimizer and NR-IQA analyses, and additional non-reference benchmark results.It also discusses Semi-UIR’s inherent advantage and provides further experiments.
A.1. Details of the Network Structure
AIM-Net restores underwater images using illumination-aware multi-scale processing and a gradient branch that supplies color, texture, edge, and structural information.
- AIM-Net combines an illumination-aware restoration branch with a gradient branch to restore color, structure, and high-quality gradient information.The gradient features are merged into the restoration branch as edge and structure priors.
- Illumination-aware Restoration Branch: The restoration branch uses a parallel multi-scale structure to integrate spatial and contextual information while maintaining edges and suppressing halo artifacts.Its backbone is a Multiscale Residual Block with streams at different resolutions.
- Illumination-aware Restoration Branch: Illumination guidance incorporates the illumination map to help the network perceive underwater color and light-source information.The illumination map reflects ambient light and scene-dependent degradation relevant to underwater image formation.
- Gradient Branch: The gradient branch uses image gradients as edge information and provides structural guidance for restoring sharp local regions.A supplementary gradient loss also promotes underwater image restoration.
- Given degraded image x and illumination map xL, AIM-Net outputs a restored clear image xout and restored gradient map gout.The network has 1.675M parameters and runs at 33.3 FPS on 256 × 256 images.
A.2. Details of Validation Set and Loss Functions
The appendix describes the validation set, loss components, optimizer choice, and NR-IQA analysis used to select reliable pseudo labels.
- Validation Set: The validation set contains 200 paired full-reference underwater images from [25], split between testS and testR at a 12:8 ratio.
- Loss Functions: Lper uses a pretrained VGG-16 network, while Lgrad constrains AIM-Net’s restored gradient map to approach the ground-truth gradient map.
- Optimizer: AdamP is selected because Adam requires longer training than AdamP to reach similar accuracy.
- NR-IQA Metrics: MUSIQ achieves the best performance among the seven evaluated NR-IQA metrics and selects more visually pleasing reliable-bank pseudo labels.
A.4. Inherent Advantage of Semi-UIR
The experiments test whether Semi-UIR benefits a general restoration architecture by replacing AIM-Net with the five-layer Unet used by FUnIEGAN.
- Experimental Setup: Replacing AIM-Net with the five-layer FUnIEGAN Unet preserves the experimental training details and datasets for a controlled comparison.
- Findings: Semi-UIR improves Unet’s generalizability on real-world underwater benchmarks and demonstrates extensibility beyond AIM-Net.
A.5. Additional Experimental Results on Non-reference Benchmark
Additional experiments compare Semi-UIR with state-of-the-art methods on non-reference underwater benchmarks and report stronger restoration quality in visual comparisons.
- Benchmark Comparisons: Semi-UIR is visually compared with GDCP, MMLE, WaterNet, Ucolor, FUnIE-GAN, PRWNet, and CWR on non-reference underwater benchmarks.
- Visual Results: Semi-UIR outperforms the compared algorithms in restoring underwater images with rich details and natural color.
- Quantitative Evaluation: Table 8 evaluates UIEB, EUVP, RUIE, and Seathru using UIQM, UCIQE, and MUSIQ, comparing Unet-base against Unet-semi.