U-shape Transformer for Underwater Image Enhancement

In plain words

Photos taken under water look washed out and off-colour, because water soaks up and scatters light. This paper tries to undo that damage with a neural network that learns from pairs of murky photos and clean reference images. Clean references were hard to come by, so the authors built a dataset called LSUI. For each real underwater photo they picked a reference from several automatic enhancers, using image-quality scores and then votes from volunteers.

The network is a U-shaped image-to-image model: the picture is squeezed down and then expanded back up. Where a plain U-Net would copy features straight across, this model passes them through two transformer blocks. One block works across colour channels, to focus on the colours that fade most. The other works across regions of the image at the bottleneck, to focus on the most damaged areas. Training combines an adversarial term with losses that compare colours in the LAB and LCH colour spaces, which the authors picked for how well they describe colour as people see it.

TL;DR

• Underwater image enhancement (UIE) network: a U-shaped generator whose skip connections are replaced by a channel-wise transformer (CMSFFT), with a spatial transformer (SGFMT) at the bottleneck §3.2.
• Trained as a GAN with an added loss that combines RGB, LAB and LCH colour terms and a perceptual term §3.3.
• Releases LSUI, 4279 real underwater image pairs with reference images, semantic maps and transmission maps §3.1.
• Author-reported: PSNR 24.16 and SSIM 0.93 on Test-L400, the best of the ten compared methods tab. 3. Parts of the released code differ from the paper, see Paper vs code.

Problem

Light absorption and scattering in water cause colour casts, colour artefacts and blur §1. Data-driven UIE was held back by two things. First, existing datasets were small, covered few scenes, or lacked matched clear references §1 §2.2. Second, most methods ignore that attenuation differs across colour channels and image regions, and CNNs focus on local features rather than long-range context §1 §2.3.

Method

The paper contributes a dataset, a generator and a loss. Equations are numbered in order of appearance, as in the compiled paper.

Dataset (LSUI) §3.1. The authors collected 8018 underwater images, some their own and some from public datasets; the paper says all were licensed for academic use only §3.1. For each image, 18 existing UIE methods produced candidate references. UIQM and UCIQE scored the candidates with equal weights, keeping the top three per image §3.1. Twenty volunteers then rated each candidate on five criteria (contrast, saturation, colour correction, artefacts, over- or under-enhancement) on a 0 to 10 scale; the scores were normalised and summed to a maximum of 100, and the top one was kept. Images whose best sum was below 70 were removed §3.1. A second round let volunteers flag defects in each reference; the flagged references were processed with three restoration methods §3.1. Another vote removed pairs that more than half of volunteers rejected §3.1. The final set has 4279 pairs, plus a hand-labelled segmentation map and a generated medium transmission map for each image §3.1 fig. 3.

Generator: U-shape Transformer §3.2.1 fig. 4. The generator follows pix2pix §1 §3.2.1. Encoding: the input is downsampled three times; each downsampled copy passes through a 1×1 convolution and then a convolution block at its scale. The outputs of four encoder blocks feed the CMSFFT and the SGFMT §3.2.1. Decoding: the SGFMT output goes through the first decoder block; four decoder blocks at different scales each receive one CMSFFT output §3.2.1.

Discriminator §3.2.1. The discriminator takes four scales. At each scale its input is the map from the layer above, the decoder map of the same size, and a 1×1 convolution of the reference image downsampled to that size §3.2.1. The paper gives no further architecture detail (not stated).

SGFMT, spatial transformer at the bottleneck §3.2.2 fig. 5. The input feature map has size H/16 × W/16 × C. Linear projection turns it into a token sequence, and a learnable position embedding is added eq. 1. Four standard transformer layers follow, each with multi-head attention and a feed-forward network with layer normalisation eq. 2 eq. 3. The last output is reshaped back into a feature map §3.2.2.

S_{in} = W * F_{in} + \mathrm{PE} eq. 1

S_{l}^{'} = \mathrm{MHA}(\mathrm{LN}(S_{l-1})) + S_{l-1} eq. 2

S_{l} = \mathrm{FFN}(\mathrm{LN}(S_{l}^{'})) + S_{l}^{'} eq. 3

CMSFFT, channel-wise transformer replacing the skip connections §3.2.3 fig. 6. The four encoder maps F_i (i = 0..3) are patch-embedded with convolutions whose kernel and stride are P/2^i, with P = 32 §3.2.3. Each scale then yields d = HW/P² tokens of C_i channels. Queries Q_i come from each scale eq. 4. Keys and values come from the channel-wise concatenation S of all four scales, with C = C_1 + C_2 + C_3 + C_4 eq. 4. Attention runs along the channel axis, not the patch axis: a channel-by-channel similarity map is normalised with instance normalisation and a softmax eq. 5. The N heads are averaged and Q_i is added back eq. 6. A feed-forward block with layer normalisation is applied, and the step repeats l = 4 times eq. 7 §3.2.3. The outputs are reshaped into feature maps for the decoder §3.2.3.

Q_i = S_i W_{Q_i} \quad K = S W_K \quad V = S W_V eq. 4

\mathrm{CA}_{i} = \mathrm{SoftMax}\left(\mathrm{IN}\left( \frac{Q_{i}^{T} K}{\sqrt{C}}\right)\right) V^{T} eq. 5

\mathrm{CMHA}_{i} = \frac{\mathrm{CA}_{i}^{1} + \cdots + \mathrm{CA}_{i}^{N}}{N} + Q_{i} eq. 6

O_{i} = \mathrm{CMHA}_{i} + \mathrm{MLP}(\mathrm{LN}(\mathrm{CMHA}_{i})) eq. 7

Colour spaces for the loss §3.3 §4.3. Generated and reference images are both converted from RGB to LAB and to LCH eq. 8 eq. 9. The authors cite the wider colour gamut and the more accurate description of saturation and brightness as the reason §3.3, and LAB's closeness to human visual characteristics §4.3. The LAB loss compares the L channel by squared difference and the A and B channels by a quantised-distribution cross-entropy eq. 10. The LCH loss uses a quantised-distribution term for L, and squared differences for C and H eq. 11. The paper calls the quantisation operator Q but does not give its bins (not stated).

L^{G(x)},A^{G(x)},B^{G(x)} = \mathrm{RGB2LAB}(G(x)), \quad L^{y},A^{y},B^{y} = \mathrm{RGB2LAB}(y) eq. 8

L^{G(x)},C^{G(x)},H^{G(x)} = \mathrm{RGB2LCH}(G(x)), \quad L^{y},C^{y},H^{y} = \mathrm{RGB2LCH}(y) eq. 9

\mathrm{Loss}_{LAB}(G(x),y) = \mathbb{E}_{x,y}\left[(L^y - L^{G(x)})^2 - \sum_{i=1}^{n} Q(A_i^y)\log Q(A_i^{G(x)}) - \sum_{i=1}^{n} Q(B_i^y)\log Q(B_i^{G(x)})\right] eq. 10

\mathrm{Loss}_{LCH}(G(x),y) = \mathbb{E}_{x,y}\left[ - \sum_{i=1}^{n} Q(L_i^y)\log Q(L_i^{G(x)}) + (C^y - C^{G(x)})^2 + (H^y - H^{G(x)})^2\right] eq. 11

The RGB term is an L2 loss and the perceptual term follows Johnson et al. §3.3. The adversarial term is the standard GAN objective eq. 12. The full generator objective is eq. 13:

L_{GAN}(G,D) = \mathbb{E}_y[\log D(y)] + \mathbb{E}_x[\log(1 - D(G(x)))] eq. 12

G^{*} = \arg\min_{G}\max_{D} L_{GAN}(G,D) + \alpha\,\mathrm{Loss}_{LAB}(G(x),y) + \beta\,\mathrm{Loss}_{LCH}(G(x),y) + \gamma\,\mathrm{Loss}_{RGB}(G(x),y) + \mu\,\mathrm{Loss}_{per}(G(x),y) eq. 13

with α = 0.001, β = 1, γ = 0.1, μ = 100, set "with numerous experiments" §3.3.

The pseudocode below follows the paper's text. It is an implementation sketch, not a line-by-line copy of the released code (see Paper vs code).

# Generator forward pass  [§3.2.1] [fig. 4]
# Input: image (B, 3, 256, 256) with values in [0, 1]  [§4.1]
# Output: four RGB maps at scales 1/8, 1/4, 1/2 and 1  [§3.2.1]
def generate(image):
    # Downsample the input three times, one copy per scale  [§3.2.1]
    scaled_inputs = downsample_three_times(image)
    # Encoder: one convolution block per scale, fed with its input copy
    skips, deepest = encode(image, scaled_inputs)
    # Channel-wise transformer replaces the skip connections  [§3.2.3]
    skips = cmsffts(skips)
    # Spatial transformer at the bottleneck  [§3.2.2]
    bottleneck = sgfmt(deepest)
    # Decoder: upsample, combine with the matching skip, convolve  [§3.2.1]
    outputs = decode(bottleneck, skips)
    return outputs

# Channel-wise attention for the four scales  [§3.2.3] [eq. 4-6]
# Tokens per scale: d = HW / P^2, with P = 32  [§3.2.3]
def channel_attention(feature_maps):
    # Patch-embed scale i with kernel and stride P / 2^i  [§3.2.3]
    seqs = [patch_embed(f, 32 // 2 ** i) for i, f in enumerate(feature_maps)]
    # Shared keys and values from the concatenated scales  [eq. 4]
    shared = concat_channels(seqs)           # (B, d, C), C = sum of C_i
    keys = shared @ W_K
    values = shared @ W_V
    outputs = []
    for i, seq in enumerate(seqs):
        # Per-scale queries  [eq. 4]
        queries = seq @ W_Q[i]
        # Channel-by-channel similarity, scaled, instance-normalised  [eq. 5]
        scores = instance_norm(queries^T @ keys / sqrt(C))
        weights = softmax(scores)            # normalised over the C axis
        context = weights @ values^T
        # Average over N = 4 heads, then add the queries back  [eq. 6]
        # Orientation of the added Q_i is not stated; see Pitfalls.
        cmha = mean_over_heads(context) + queries
        # Feed-forward step, repeated l = 4 times  [eq. 7] [§3.2.3]
        outputs.append(cmha + mlp(layer_norm(cmha)))
    return outputs

# Generator objective  [eq. 13]
# output and reference are RGB tensors; score is the discriminator output
def generator_loss(output, reference, score):
    # Adversarial term  [eq. 12]
    adversarial = gan_term(score)
    # L2 on RGB, weight gamma = 0.1  [§3.3] [§4.1]
    rgb = l2(output, reference)
    # Perceptual loss in a VGG-19 feature space, weight mu = 100  [§3.3]
    perceptual = vgg_distance(output, reference)
    # LAB and LCH terms, weights alpha = 0.001 and beta = 1  [eq. 10] [eq. 11]
    lab = lab_loss(output, reference)
    lch = lch_loss(output, reference)
    return (adversarial + 0.001 * lab + 1.0 * lch
            + 0.1 * rgb + 100.0 * perceptual)

Hyperparameters

Values as stated in the paper. Where the released code uses a different value, the code value is shown in the notes (details in Paper vs code).

Results

Author-reported, on the paper's own splits. Setups are given per table. Test-L400 and Test-U90 are full-reference sets (they have clean references). Test-U60 and SQUID are non-reference sets §4.2.

Colour space for the loss §4.3 tab. 1. PSNR of U-shape Transformers trained with a single extra colour-space loss, trained on Train-L. The top three per test set are LAB, RGB and LCH (Test-L400: LAB 23.86, RGB 23.79, LCH 23.62; Test-U90: RGB 22.72, LAB 22.53, LCH 22.49).

Dataset evaluation §4.4 tab. 2. The same network retrained on different training sets. The LSUI-trained model scores best on both metrics. The source marks the top scores in red.

Full-reference comparison §4.5 tab. 3 fig. 8. Test-L400 and Test-U90. Timing is at 256 × 256 on the hardware in the table caption. "×" is printed in the source. Its meaning is not stated; it is probably "not applicable" for the physical and prior-based methods.

Non-reference comparison §4.5 tab. 4 fig. 9. PS is the perception score (1 to 5, higher is better). UIQM and UCIQE are higher-is-better; NIQE is lower-is-better. The source marks the best value of each column in red; the marking is not reproduced here.

Colour restoration §4.6 tab. 5 fig. 10. CIEDE2000 colour difference on Color-Checker7 (7 images, lower is better). The source marks the best per column in red; the average column is as printed.

Ablation §4.7 tab. 6 fig. 11. BL is the pix2pix baseline §4.7. Test-L400 and Test-U90, all trained on Train-L.

Headline reading. On Test-L400 the full model is 1.25 dB PSNR above the next best method, Ucolor (24.16 vs 22.91, from tab. 3). On Test-U90 it is 2.13 dB above Ucolor (22.91 vs 20.78). These gaps are computed here from the table values. The paper's abstract says "more than 2dB superiority" [abstract]. Pitfalls explains why that claim needs care.

Paper vs code

The search order was: the paper text, which gives no repository link; then the authors' project page, which links the official repository [github.com/LintaoPeng/U-shape_Transformer]. Its owner handle matches the first author, and the code headers name Lintao Peng as author Ushape_Trans.py. The repository is MIT-licensed. Commit 41f9897669d2e6f6718e00772ece2523d0f149a2 is the one read here. Code paths below are relative to that repository. Notebook line numbers refer to the .ipynb JSON.

The items below come from reading the code against the paper. They are differences the implementer has to resolve, not claims that the released code is wrong.

• Channel widths. Paper: C_1..C_4 = 64, 128, 256, 512 §3.2.3. The generator calls the channel transformer with [32, 64, 128, 256] Ushape_Trans.py. The encoder blocks produce 32, 64, 128 and 256 channels Ushape_Trans.py. The class default is 64..512 CMSFFT.py, which the generator overrides.
• Attention scaling. Paper: divide by sqrt(C), with C = C_1 + ... + C_4 eq. 5. Code: divides by sqrt(KV_size), with KV_size = 480 = 32 + 64 + 128 + 256 CMSFFT.py CMSFFT.py.
• Residual in the channel attention. Paper: adds Q_i back eq. 6. Code: applies an output linear layer to the averaged heads CMSFFT.py, then adds the block input, not Q_i CMSFFT.py.
• Feed-forward width. Paper: gives no expansion factor eq. 7. Code: expands the hidden size by 4 with GELU CMSFFT.py CMSFFT.py.
• Bottleneck projection. Paper: linear projection with learned position embedding §3.2.2 eq. 1. Code: BatchNorm and ReLU, then a 3×3 convolution from 256 to 512 channels Ushape_Trans.py Ushape_Trans.py, a learned position embedding of fixed size (1, 256, 512) PositionalEncoding.py, and a residual addition after the transformer Ushape_Trans.py.
• Bottleneck feed-forward width. Paper: not stated. Code: the generator passes a hidden size of 256 Ushape_Trans.py Ushape_Trans.py. A comment in the SGFMT file says 4096 SGFMT.py. The value used is 256.
• Skip gating. Paper: the CMSFFT replaces the skip connections §3.2.3. Code: each decoder skip is also gated by a CCA module before concatenation Ushape_Trans.py Ushape_Trans.py. CCA is defined in block.py. The paper does not mention CCA.
• Input downsampling. Paper: "downsampled three times" §3.2.1. Code: 2×2 max-pooling Ushape_Trans.py.
• Adversarial loss. Paper: log form eq. 12. Code: least-squares MSE with all-ones targets for real and all-zeros for fake train.ipynb train.ipynb train.ipynb.
• Extra SSIM term. Code adds −100 × SSIM over four scales to the generator loss train.ipynb train.ipynb. Eq. 13 has no SSIM term.
• Pixel loss type and phases. Paper: L2 for the first 600 epochs, L1 for the last 200 §4.1. Code: a 500-epoch run with MSE train.ipynb, then a resumed run with L1 from epoch 495 to 800 trainL1.ipynb trainL1.ipynb. The resumed run loads generator and discriminator weights from the epoch-495 checkpoint trainL1.ipynb. It recreates the optimisers and schedulers trainL1.ipynb.
• Learning-rate schedule. Paper: 5e-4, then 2e-4 after epoch 600, with × 0.8 every 40 epochs §4.1. Code: 5e-4 with step 40 train.ipynb train.ipynb; 2e-4 with step 50 in the resumed run trainL1.ipynb trainL1.ipynb.
• Batch size. Paper: 6 §4.1. Code: 1 train.ipynb.
• Multi-scale losses. Paper: no per-scale averaging stated. Code: the pixel, SSIM, VGG, LAB and LCH terms are computed at four scales and divided by 4 train.ipynb. Only the discriminator is multi-scale in the paper's text.
• LAB loss. Paper: squared L difference eq. 10. Code: absolute L difference LAB.py. The AB term is a cross-entropy over a joint 7 × 7 grid on [−80, 80] with a softmax on squared distance LAB.py LAB.py LAB.py. The paper gives no bin count.
• LCH loss. Paper: squared C and H differences eq. 11. Code: absolute L and C differences LCH.py, and a 4-bin soft-histogram cross-entropy for H LCH.py LCH.py. The reference hue distribution is built with eps = 0.0, while the prediction uses eps = 0.01 LCH.py.
• Reduction. Paper: no reduction stated. Code: the MSE and L1 criteria sum over pixels train.ipynb trainL1.ipynb. The AB cross-entropy also sums over pixels LAB.py, while its L term is a mean LAB.py. The effective weights in eq. 13 depend on these choices.
• Perceptual loss. Paper: no layer or input normalisation stated §3.3. Code: torchvision VGG-19 with ImageNet weights utils.py. The feature is taken at module index '30', labelled conv5_2 utils.py. The loss is a mean squared difference utils.py. No mean and std normalisation is visible in that module utils.py.
• Augmentation. Paper: cropping, rotation and flipping §4.1. The augmentation helper exists utils.py, but neither training notebook calls it.
• Output range. The generator's last layer is a linear 1×1 convolution with no activation block.py. Training clamps to [0, 1] only for the logged PSNR train.ipynb. The test notebook saves the output with torchvision's save_image(..., normalize=True) test.ipynb, which rescales each output by its own minimum and maximum before writing it. The saved images are therefore not the raw network output.
• Released weights. The test notebook loads epoch-795 generator weights test.ipynb. The README links the pretrained models.
• Discriminator. Paper: no architecture §3.2.1. Code: patch output (1, 8, 8) train.ipynb. Full definition: Ushape_Trans.py.

Pitfalls

• LSUI licence. The paper says the collected images were "licensed and used only for academic purposes" §3.1. The project page has no licence statement; the MIT licence on the repository covers the code, not the dataset LICENSE. The dataset's licence is therefore not stated. Do not redistribute or use LSUI in commercial training without checking first.
• Dataset size. The arXiv abstract in our metadata says 5004 image pairs. The LaTeX abstract and body say 4279 §3.1. A commented-out line in the source still says 5004. Use 4279 and note the discrepancy.
• Split. The paper says the 3879 / 400 split was random §4.1. The repository says to draw it randomly, and links an author-made split [README]. Which split the reported numbers used is not stated.
• "More than 2 dB". The abstract claims more than 2 dB superiority. From tab. 3, that holds against Ucolor on Test-U90 (2.13 dB) but not on Test-L400 (1.25 dB). It also holds only where the best competitor is Ucolor.
• Efficiency claim. The results text says "relatively few parameters, FLOPs and running time" §4.5. In tab. 3, the model's 66.2 GFLOPs and 65.6 M parameters are higher than FUnIE (10.23 G, 7.019 M) and UIE-DAL (29.32 G, 18.82 M). It has fewer FLOPs than WaterNet (193.7 G) and Ucolor (443.85 G), but more parameters than WaterNet (24.81 M). Its runtime (0.07 s) is below WaterNet, Ucolor and FUnIE, and UGAN is faster (0.05 s).
• Equation numbers. The source labels the GAN and total objectives eq:14 and eq:15. Counted in source order, the displayed numbers are 12 and 13, and this card uses those. Check against the compiled PDF before quoting a number.
• Table 5 average. The UIBLA average is printed as "15,15" in the source tab. 5. Read it as 15.15.
• Perception survey. The body text does not state how many raters scored the PS metric. An older commented-out sentence in the source says 10 volunteers. Treat the count as not stated.
• Ambiguous CMSFFT shapes. The paper writes Q_i as d × C_i and W_Q as d × C_i eq. 4. The text gives CA_i as C_i × d §3.2.3, yet eq. 6 adds Q_i to it. The orientation of the added term is therefore unclear. The code adds the block input instead CMSFFT.py.
• Quantisation bins. The quantisation operator Q is never given its bin count or range eq. 10 eq. 11. The code uses 7 bins on [−80, 80] for AB and 4 hue bins LAB.py LCH.py. Treat both as code-only choices.
• Loss scales. The weights in eq. 13 are only meaningful with the reduction the code uses (see Paper vs code). The LAB cross-entropy is summed over every pixel LAB.py, so with a per-pixel mean the same α = 0.001 would weigh that term about 65 000 times less at 256 × 256. Keep the sum if you copy the weights.
• Hue wrap-around. The released LCH code wraps the hue bins: bin 0 also receives the mass of the last bin LCH.py. The paper does not describe this.
• Unbounded output. The network has no output activation. An implementation that clamps the output or applies a sigmoid changes what the network learns to output.
• Test scaling. Test images are resized to 256 × 256 and outputs are min-max normalised when saved test.ipynb. Compute PSNR on the raw output, not on the saved PNG.
• Run time. The paper gives no training time. The released notebook's only log is an interrupted first epoch with an estimated time remaining of about 1 day 19 hours train.ipynb. Hardware and run length differ, so treat this as one data point, not a benchmark.

Lineage

• Builds on: pix2pix (Isola et al., cited as [isola2017image]). The generator and discriminator are based on it §1 §3.2.1 §4.7.
• Channel attention idea: adapted from Wang et al., cited as [wang2021uctransnet], for the CMSFFT §3.2.3. The code header says "Design based on the CTrans" CMSFFT.py.
• Transformer blocks: ViT [Dosovitskiy2021AnII] and the original transformer [vaswani2017attention] §2.3 §3.2.2.
• Compared against: UIBLA, UDCP, Fusion, Retinex, RGHS, WaterNet, FUnIE, UGAN, UIE-DAL and Ucolor §4.2. Ucolor is the strongest competitor on the full-reference sets tab. 3.
• Superseded by: not known from the content of this wiki.