Dynamic SpectraFormer for Ultra-High-Definition Underwater Image Enhancement

AI summary

Overview

• Research area: Computer vision, specifically underwater image enhancement (UIE) using frequency-domain transformers, aimed at ultra-high-definition (UHD) imagery.
• Technical level: Advanced. The paper assumes familiarity with transformer self-attention, the Discrete Cosine Transform (DCT), and complexity analysis of attention mechanisms.
• Scope in one sentence: The paper proposes a frequency-domain transformer, Dynamic SpectraFormer, that performs sparse attention over DCT spectrum bands selected adaptively per input image, and reports state-of-the-art PSNR/SSIM on three underwater benchmarks (UIEB, LSUI, EUVP) at lower FLOPs than competing transformer and CNN baselines.

What This Paper Is About

Underwater photos degrade because light is absorbed and scattered in water: red light (longest wavelength) disappears first, then green, then blue, which produces color casts, haze, and poor visibility that hamper marine robots and Autonomous Underwater Vehicles (AUVs). The authors observe that color and brightness distortions live mostly in low frequencies, while edge and texture distortions live mostly in high frequencies, so an effective method must fix both at once. Their goal is a transformer that enhances UHD underwater images in the frequency domain, adapts its frequency weighting to each input, and stays computationally feasible at high resolutions where standard pixel-level self-attention is quadratic in spatial resolution, on the order of O(W²H²).

Key Contributions

1. A sparse spectrum attention block (Dynamic SpectraFormer). Query, key, and value matrices are transformed with the DCT, attention is computed over a sparsified set of frequency bands, and an inverse DCT (iDCT) returns the attended features to the spatial domain. The paper claims this retains global attention and long-range dependency modeling while reducing attention cost to O(HWKC) instead of the O(H²W²C) of pixel-level attention.
2. A Dynamic Spectrum Weight Generator (DSWG). A lightweight module that acts as an adaptive frequency-band selector, amplifying informative bands and suppressing irrelevant ones. It uses a zigzag scan of the DCT spectrum, average-pooling downsampling to K bands, a two-layer fully connected design with LayerNorm and GELU, softmax, then inverse zigzag and upsampling to produce per-input spectrum weights M.
3. A multi-scale U-shaped encoder–decoder for UHD restoration. A 4-level symmetric encoder–decoder with Dynamic SpectraFormer blocks, 3×3 convolutions for low-level features, pixel-unshuffle/pixel-shuffle downsampling and upsampling, and skip connections, trained with a hybrid loss combining an L1 pixel loss and MS-SSIM with weights w1 = 0.6 and w2 = 0.4.
4. Strong empirical results plus released code. The authors report the best performance on all three evaluation sets and publish the source code at https://github.com/arifence2024/DynamicSpectraFormer.git.

Main Findings

• Best reported accuracy across all three benchmarks. Dynamic SpectraFormer reaches 25.98 dB PSNR / 0.9341 SSIM on UIEB, 26.33 dB / 0.9327 on LSUI, and 29.78 dB / 0.8848 on EUVP, with 64.2M parameters and 50.2 GFLOPs (FLOPs measured at 256×256 input).
• Improvement over the strongest UHD baseline. On UIEB the method improves PSNR by 0.94 dB over UHD Underwater Enhancement (25.04 dB, 0.9158 SSIM, 157M parameters, 70.23 GFLOPs).
• Improvement over the transformer baseline at lower cost. It gains 3.07 dB PSNR on UIEB over U-Shape Transformer (22.91 dB, 0.9100 SSIM; 24.16 dB, 0.9322 SSIM on LSUI; 65.6M parameters; 70.2 GFLOPs), while using fewer FLOPs (50.2 GFLOPs vs. 70.2 GFLOPs).
• Other baselines are lower. As reported in Table I: Ucolor 20.78 dB / 0.8713 SSIM (UIEB) and 22.91 dB / 0.8902 (LSUI); WaterNet 19.81 dB / 0.8612 (UIEB), 17.73 dB / 0.8223 (LSUI), 20.14 dB (EUVP); UGAN 20.68 dB / 0.8430 (UIEB), 19.79 dB / 0.7843 (LSUI), 23.49 dB (EUVP); FUnIE-GAN 19.45 dB / 0.8602 (UIEB), 19.37 dB / 0.8401 (LSUI), 23.40 dB / 0.8420 (EUVP); Deep SESR is reported only for EUVP at 24.21 dB / 0.8401 SSIM.
• The DSWG is the key component. Replacing it with randomly selected spectrum bands collapses UIEB PSNR from 25.98 dB to 21.01 dB (SSIM 0.8402), while an all-pass filter (all weights set to 1) reaches 23.72 dB / 0.8461 on UIEB, 25.20 dB / 0.8753 on LSUI, and 28.09 dB / 0.8654 on EUVP. The full DSWG gives the best scores on all three datasets.
• Sparse spectra are sufficient. With only 32 spectrum bands the model reaches 24.23 dB PSNR on UIEB (0.9212 SSIM), and increasing the length from 64 to 256 leaves accuracy nearly unchanged (128 bands: 26.00 dB UIEB, 26.33 dB LSUI, 29.82 dB EUVP; 256 bands: 25.97 dB UIEB, 26.19 dB LSUI, 29.64 dB EUVP). The authors therefore set top-K = 64. Very short spectra hurt: top-K = 8 yields 19.45 dB UIEB, 19.71 dB LSUI, 19.32 dB EUVP.
• Adaptive weighting follows image content. In the LSUI examples, images with rich detail produce dynamic spectrum weights concentrated on high-frequency bands, while plain images emphasize low-frequency bands; lighter color in the visualization indicates larger spectrum weight.
• Complexity is reduced by construction. Sparse attention is O(HWKC) versus O(H²W²C) for pixel-level attention; per-channel DCT/iDCT cost is O(HWC⌈log₂(HW)⌉) with the fast DCT algorithm; DSWG cost is O(KV²C + K²VC). Fig. 7 compares FLOPs and memory against ViT and U-Shape Transformer across image resolutions.

Methodology in Plain English

The authors start from a simple visual observation: the DCT spectrum of an enhanced underwater image has more high-frequency energy than the degraded version, so the enhancement problem can be attacked in the frequency domain. They build a U-shaped encoder–decoder network in which the core block replaces spatial self-attention with attention over frequency coefficients.

Inside each block, 1×1 convolutions produce query, key, and value feature maps. Each is converted with a 2D DCT, which is real-valued and therefore cheaper than a Fourier transform for a neural network. A small gating module, the DSWG, looks at the spectrum of each query (flattened by a zigzag scan from low to high frequency, pooled down to a fixed number of bands, passed through a two-layer fully connected network with LayerNorm and GELU, then softmax) and produces weights that re-scale the frequency bands. Only the top-K bands after this weighting are kept for computing the attention map, which is what makes the computation tractable at UHD resolutions. Attention is computed over these sparse frequency components, the result is mapped back to the spatial domain with an inverse DCT, and the outputs from all heads are concatenated and passed through a linear normalization layer.

The network is trained on patches resized to 512×512 with flipping, noise, and contrast augmentation, using the Adam optimizer with initial learning rate 3×10⁻⁴ and cosine annealing, implemented in PyTorch on 4× NVIDIA Tesla A100 GPUs. Supervision combines an L1 pixel loss (weight 0.6) with MS-SSIM (weight 0.4). Training data follows the U-Shape Transformer setup: LSUI split into 4500 training and 404 testing images, with evaluation on UIEB (90 pairs), LSUI (504 pairs), and EUVP (515 pairs) using PSNR and SSIM.

Why This Matters

Impact on research. The paper challenges the common practice in frequency-domain image restoration of applying an image-agnostic global filter to the whole spectrum, arguing that treating all frequency bands equally is suboptimal and that content-dependent band selection is better. It also offers a concrete route to global attention in transformers without the O(W²H²) blow-up that forces prior UHD work into local windows — a trade-off the authors argue causes mixed-up details and inaccurate hues in restored images. The reported gains (0.94 dB over the previous best UHD method on UIEB at fewer FLOPs and fewer parameters than several baselines) give the community a new reference point for frequency-domain UHD restoration.

Real-world applications (as motivated by the paper):

• Visual sensing for Autonomous Underwater Vehicles (AUVs) and Remotely Operated Vehicles (ROVs), which the paper says is hampered by degraded imagery despite high-end cameras.
• Marine robots and aquatic exploration more broadly, where clear imagery supports perception and navigation.
• On-device operation with

Authors’ abstract

Underwater images suffer from color distortion, haze, and poor visibility due to light refraction and absorption in water. These challenges significantly impact the utilization of Autonomous Underwater Vehicles (AUVs) or marine robots. Typically, color and brightness distortions manifest at lower frequencies, while edge and texture distortions are prevalent at higher frequencies. Traditional methods struggle to concurrently rectify these mixed distortions as they primarily concentrate on the spatial domain. To address these issues, we introduce the Dynamic SpectraFormer, which enhances underwater images through a frequency domain transformer. The Dynamic SpectraFormer introduces an ultra-high-resolution sparse spectrum attention module, which could capture the long-term dependency without losing the universal approximating power. Additionally, we have developed a dynamic spectrum weight generation layer that serves as an adaptive spectrum band selector, accentuating critical frequency bands and suppressing less relevant ones. Consequently, this method significantly improves underwater image quality by addressing both high- and low-frequency distortions. Our extensive ablation studies and comparative evaluations consolidate the Dynamic SpectraFormer's efficacy across multiple underwater image enhancement benchmarks. The source code is available at https://github.com/arifence2024/DynamicSpectraFormer.git.