PaperlayerView exact revision on arXiv

Source-linked AI summary

U-shape Transformer for Underwater Image Enhancement

Lintao Peng, Chunli Zhu, Liheng Bian

arXiv:2111.11843v6cs.CVeess.IV

TL;DR

Underwater image enhancement is limited by degraded imaging, insufficient large-scale paired datasets, and incomplete modeling of attenuation across color channels and spatial regions. The paper introduces the LSUI dataset and a U-shape Transformer with channel-wise and spatial-wise transformer modules plus a multi-color-space loss. Extensive experiments validate state-of-the-art UIE performance, while ablations report improved quantitative and visual results from combining the proposed modules.

  • Problem

    Underwater enhancement research lacks large-scale paired data with varied scenes and high-quality references, while CNN-based methods inadequately characterize channel- and region-dependent attenuation and global dependencies.

  • Method

    The paper constructs the LSUI dataset and proposes a U-shape Transformer integrating CMSFFT, SGFMT, and a loss combining RGB, LAB, and LCH color spaces.

  • Results

    The full model achieves the best quantitative performance on two testing datasets, while full-reference evaluation reports the best PSNR and SSIM with relatively few parameters, FLOPs, and running time.

  • Takeaways & Limitations

    Combining channel-wise and spatial-wise attention with multi-color-space supervision improves enhancement quality across quantitative and visual evaluations.

Abstract

from arXiv · show

The light absorption and scattering of underwater impurities lead to poor underwater imaging quality. The existing data-driven based underwater image enhancement (UIE) techniques suffer from the lack of a large-scale dataset containing various underwater scenes and high-fidelity reference images. Besides, the inconsistent attenuation in different color channels and space areas is not fully considered for boosted enhancement. In this work, we constructed a large-scale underwater image (LSUI) dataset including 5004 image pairs, and reported an U-shape Transformer network where the transformer model is for the first time introduced to the UIE task. The U-shape Transformer is integrated with a channel-wise multi-scale feature fusion transformer (CMSFFT) module and a spatial-wise global feature modeling transformer (SGFMT) module, which reinforce the network's attention to the color channels and space areas with more serious attenuation. Meanwhile, in order to further improve the contrast and saturation, a novel loss function combining RGB, LAB and LCH color spaces is designed following the human vision principle. The extensive experiments on available datasets validate the state-of-the-art performance of the reported technique with more than 2dB superiority.

I. INTRODUCTION

Underwater image enhancement is important but difficult because absorption and scattering produce color and detail degradation, while existing methods do not fully model underwater-specific attenuation or consistently paired training data. The paper addresses these gaps with the LSUI dataset, a U-shape Transformer, and a multi-color-space loss.

  • Motivation: Light absorption and scattering cause underwater images to exhibit color casts, color artifacts, and blurred details.These degradations affect applications including ocean exploration, biology, archaeology, and underwater robotics.
  • Existing methods: Existing UIE methods include physical-model, visual-prior, and data-driven approaches, each with limitations in complex underwater scenes.Visual priors neglect physical degradation, physical models depend on suitable scene assumptions, and data-driven methods lack well-matched real-world training pairs.
  • Results: The reported method achieves the highest PSNR and best visual quality in the cited comparison and outperforms evaluated methods on PSNR and SSIM.Full-reference evaluation reports the best PSNR and SSIM performance with relatively few parameters, FLOPs, and running time.
  • Contributions: The LSUI dataset contains 4279 real underwater images with generated clear references, semantic segmentation maps, and medium transmission maps.It covers more abundant water types, lighting conditions, and target categories than existing underwater datasets.
  • Contributions: The U-shape Transformer uses channel-wise and spatial-wise attention modules to address inconsistent attenuation across color channels and spatial regions.CMSFFT and SGFMT are embedded in the network to focus on areas and channels with more serious attenuation.
  • Contributions: The multi-color-space loss combines RGB, LAB, and LCH features to improve output contrast and saturation.The design follows the paper’s stated human-vision-oriented color-space selection.

B. Underwater Image Datasets

Existing underwater image datasets are difficult to construct because real-world scenes are dynamic and matched training data are hard to collect. Available datasets include non-reference datasets and other dataset types described in the paper.

  • Dataset motivation: The sophisticated and dynamic underwater environment makes collecting matched underwater training data extremely difficult.This motivates dataset designs that represent varied underwater lighting, depth of field, blurriness, and color casts.

C. Transformers

The paper motivates transformers for UIE because CNNs struggle with inconsistent attenuation across channels and spatial regions and with global dependencies. It also constructs LSUI with diverse real-world scenes, generated references, and auxiliary maps.

  • Motivation: CNN-based UIE methods struggle to characterize inconsistent attenuation across color channels and spatial regions and to model long-range global features.Transformers are introduced as suitable because self-attention models long-distance dependencies and global features.
  • Dataset: The dataset construction selects diverse real underwater images and uses generated references, metric scoring, volunteer ratings, and subsequent optimization.Reference-image refinement addresses blur, color cast, and noise before final dataset assembly.
  • Dataset: Each LSUI image has a corresponding clear reference image, semantic segmentation map, and medium transmission map.These annotations support both comparison and further UIE development.
  • Dataset: LSUI contains 4279 real-world underwater images with richer water types, lighting conditions, and target categories than existing datasets.The dataset includes deep-sea and underwater cave scenes unavailable in previous datasets.

B. U-shape Transformer

The U-shape Transformer combines a CMSFFT- and SGFMT-based generator with a discriminator and multiscale connections. SGFMT replaces the bottleneck to model global information, while transformer layers process projected feature sequences.

  • Overall Architecture: The U-shape Transformer uses a CMSFFT- and SGFMT-based generator together with a discriminator.Its encoder produces multiscale features for the transformer modules, and the decoder reconstructs enhanced images.
  • Overall Architecture: Multiscale connections link generator and discriminator features, allowing freer gradient flow and more stable training.The discriminator receives features from its own upper layer, the generator decoder, and downsampled reference images.
  • SGFMT: SGFMT replaces the generator bottleneck and uses spatial self-attention to model global information and emphasize severely degraded regions.The input feature map is projected into a one-dimensional sequence with learnable positional embeddings.
  • SGFMT: The SGFMT transformer block contains four standard transformer layers, each combining multi-head attention and a feed-forward network.The final sequence is remapped into a feature map for subsequent decoding.

3) CMSFFT:

CMSFFT replaces the generator skip connection with channel-wise multi-scale feature fusion. It projects encoder features at four scales, applies channel-wise attention, and remaps the outputs for decoder blocks.

  • CMSFFT: CMSFFT replaces the original generator skip connection to emphasize color channels with more severe attenuation.It performs channel-wise multi-scale fusion on encoder feature maps before transmitting results to the decoder.
  • Multi-Scale Feature Encoding: Convolutional projections generate four feature sequences from encoder maps at different scales while preserving their channel counts.The four sequences are concatenated through the channel dimension to form shared keys and values.
  • Channel-Wise Attention: Channel-wise multi-head attention operates along the channel axis rather than the classical patch axis.Instance normalization on similarity maps assists gradient flow while guiding attention toward more degraded channels.
  • Transformer Fusion: The attention outputs are averaged across four heads, processed by a feed-forward network, and repeated four times to build the transformer.Feature remapping reorganizes the outputs into four multiscale feature maps for the decoder.

C. Loss Function

The proposed loss combines RGB, LAB, and LCH representations with perceptual and adversarial objectives. LAB and LCH broaden color representation and describe saturation and brightness more accurately for training.

  • Multi-Color Space Loss: The loss function combines RGB, LAB, and LCH color spaces to improve color-related supervision.The generator output is converted from RGB into LAB and LCH before computing color-space losses.
  • Multi-Color Space Loss: LAB and LCH are used because they provide wider color gamut representation and more accurate descriptions of color saturation and brightness.The loss compares generated output with the original input and reference image in these color spaces.
  • Optimization Objective: The generator objective includes RGB L2 loss, perceptual loss, LAB loss, and LCH loss.These four losses are combined with the adversarial objective.
  • Optimization Objective: The final loss weights α, β, γ, and µ are set to 0.001, 1, 0.1, and 100, respectively.The paper reports that these hyperparameters were selected through numerous experiments.

IV. EXPERIMENTS

The experiments evaluate training details, color-space selection, dataset quality, state-of-the-art comparisons, and component effectiveness across multiple underwater datasets.

  • Experiment scope: Experiments assess color-space selection, dataset quality, state-of-the-art UIE performance, and U-shape Transformer components.The study evaluates five datasets and includes ablation studies for each network component.

A. Implementation Details

The LSUI data are split into training and testing subsets, augmented and normalized before network input, while training uses scheduled Adam optimization over 800 epochs.

  • Data preparation: LSUI is divided into Train-L with 3879 images and Test-L400 with 400 images.Training images are augmented through cropping, rotation, and flipping.
  • Data preparation: Input images are resized to 256*256 and normalized to [0,1].
  • Optimization: The model trains for 800 epochs with Adam, batch size 6, and learning rates of 0.0005 then 0.0002 across two training phases.The learning rate decreases 20% every 40 epochs.
  • Optimization: LossRGB uses L2 loss during the first 600 epochs and L1 loss during the final 200 epochs.

B. Experiment Settings

The experiments compare training datasets and evaluate full-reference and non-reference performance, including color-space loss choices and perceptual metrics.

  • Training and testing datasets: Train-U combines 800 UIEB image pairs with 1,250 synthetic underwater images, while Train-E uses paired EUVP training images.
  • Training and testing datasets: Testing uses full-reference Test-L400 and Test-U90 datasets alongside non-reference Test-U60 and SQUID datasets.
  • Color-space selection: Color-space experiments test U-shape Transformers with different loss functions on Test-L400 and Test-U90, ranking spaces by PSNR.The top three PSNR color spaces are marked in the table.
  • Evaluation metrics: Perception score ranges from 1-5 with higher values indicating better quality, whereas lower NIQE indicates higher visual quality.
  • Color-space selection: LAB, LCH, and RGB achieve the top-three PSNR scores on both test datasets and are combined into the multi-color-space loss.LAB aligns more closely with human visual characteristics, while LCH expresses brightness, saturation, and hue.

D. Dataset Evaluation

Dataset evaluation shows that LSUI-trained models achieve the strongest reported full-reference results and visually align closely with reference images, while the proposed method also performs strongly on non-reference metrics.

  • Dataset evaluation: LSUI-trained models achieve the best PSNR and SSIM results on Test-L400 and Test-U90.The evaluation retrains U-net, UGAN, and U-shape Transformer on Train-L, Train-U, and Train-E before testing.
  • Dataset evaluation: Train-L enhancement results have the highest PSNR and preferable visual quality, while other training datasets show color casts.LSUI contains richer lighting conditions, water types, target categories, and higher-quality reference images.
  • Method comparison: FUnIE is fast and lightweight but has limited scalability on complex and distorted samples, while other methods omit or incompletely model underwater attenuation.The comparison identifies limitations involving inconsistent characteristics, transmission priors, and multi-color-space integration.
  • Full-reference comparison: The proposed method produces results closest to the reference image, with fewer color artifacts and high-fidelity object areas.
  • Full-reference comparison: Compared methods exhibit color casts, artifacts, yellowish or reddish tones, and local over-enhancement across the visual comparisons.UIBLA shows severe color casts, Retinex improves contrast but retains casts, and UGAN and Ucolor can over-enhance locally.
  • Non-reference comparison: The proposed method achieves the highest PS and NIQE scores on the non-reference evaluation.The authors note that UCIQE and UIQM are not absolute justifications because they are insensitive to color artifacts and casts and biased toward some features.
  • Non-reference comparison: Its enhancement results have the highest PS value, while competing methods show undesirable artifacts, over-saturation, and unnatural color casts.UGAN and UIE-DAL are affected by local over-enhancement and color artifacts linked to inconsistent attenuation across areas and channels.

F. Color Restoration Performance Evaluation

The evaluation compares underwater color correction across multiple methods using Color-Checker7 and non-reference visual assessment. The U-shape Transformer achieves the lowest color dissimilarity for six cameras and the best average score.

  • Evaluation protocol: Color-Checker7 evaluates robustness across seven underwater images captured with different cameras and containing photographed color checkers.CIEDE2000 measures differences between ground-truth Macbeth Color Checker patches and enhancement results.
  • Visual comparison: The non-reference comparison presents raw images and seven enhancement methods, with perception scores shown for each result.The highest perception score in each row is marked in yellow.
  • Quantitative results: The U-shape Transformer obtains the lowest color dissimilarity for six cameras and the best average score.The reported cameras are Pentax W60, Pentax W80, Cannon D10, Fuji Z33, Panasonic TS1, and Olympus T6000.
  • Visual comparison: On the Fuji Z33 example, the U-shape Transformer achieves the highest CIEDE 2000 score among the compared UIE methods.The comparison also describes bluish, distorted, low-saturation, and reddish results among several competing methods.

G. Ablation Study

The ablation study evaluates four components of the U-shape Transformer on Test-L400 and Test-U90. The complete model performs best quantitatively, while individual modules contribute distinct visual improvements.

  • Study design: The ablation study tests CMSFFT, SGFMT, MSG, and MCSL using models trained on Train-L.The baseline model is denoted BL, while the full model is the complete U-shape Transformer.
  • Overall ablation results: The full model achieves the best quantitative performance on both Test-L400 and Test-U90.Figure 11 additionally reports the full model's highest PSNR and best visual quality.
  • Component effects: MSG reduces noise and artifacts by helping reconstruct local details in the BL+MSG result.This comparison attributes the observed visual difference to the multi-scale gradient flow mechanism.
  • Component effects: MCSL produces overall colors close to the reference, whereas CMSFFT and SGFMT improve distribution but do not by themselves produce accurate overall color.BL+MCSL still shows uneven visualization and local artifacts, while the attention modules provide more even distributions.
Loading 2111.11843v6…