Underwater Ranker: Learn Which Is Better and How to Be Better

In plain words

Judging whether an underwater photo looks good is hard. Water tints and blurs images in different ways depending on the scene, so simple formulas often disagree with what people see. This paper takes a different route. Instead of asking a model for an absolute score, it shows the model several enhanced versions of the same underwater scene, already put in order by people, and trains it to get that order right.

The model is a Transformer that looks at the image at several sizes. It also gets a summary of the colour distribution, which helps it judge how strongly the water has tinted the picture. Trained this way, its ranking agrees with human rankings better than the older hand-made formulas it was compared with.

The authors then use this learned judge to train an image-enhancement network: its opinion is added to the usual pixel-matching loss, which improves the enhanced images. They also add a small normalization step at the end of a simple enhancer, which stretches values that spill outside the valid range back into it and improves results without adding parameters.

TL;DR

URanker is a conv-attention Transformer that scores underwater images so the scores follow a human ranking of enhanced versions of the same scene. It is trained with a margin ranking loss on URankerSet, a set of 890 ranked groups. Its score, used as an extra loss, improves three enhancement networks in PSNR (author-reported). The official code is licensed CC BY-NC 4.0, the image data are not in the repository, and several settings in the code differ from the paper.

Problem

• Hand-crafted no-reference metrics (UCIQE, UIQM) do not reflect quality reliably across diverse underwater scenes and degradations §1 fig. 1.
• Full-reference metrics (PSNR, SSIM) need a reference image. The UIEB reference is one enhancement chosen by hand from several, and it is not always the best one §1.
• Fitting absolute scores is harder to learn than a relative order, and the order also avoids the scene-to-scene variance of hand-crafted metrics §1.
• Twice-Mixing generates intermediate images by blending an image with its reference. The authors argue this cannot cover over-enhanced results §2.

Method

Section numbers §n count sections in source order. The LaTeX sets secnumdepth to 0, so the compiled paper prints no section numbers; subsections are numbered here as 4.1 and so on. Equation, table and figure numbers are as printed.

Dataset (URankerSet). For each of 890 raw underwater images, nine UIE outputs are added, so each group holds ten images §3. The nine algorithms come from the prototype of Li et al. (twelve algorithms; three dropped because their outputs were similar to a kept one) §3. Eleven volunteers compare images pairwise. Each group is sorted by bubble sort in descending quality, and the sort repeats until a full pass makes no swap §3. Groups are split 800 for training and 90 for testing §6.1.

URanker, input stage. The network is a stack of serial stages, each built from a conv-attention Transformer block §4. Each stage downsamples its feature map by a patch embedding (stated as 2× per serial block) and flattens it into tokens §4. Every stage prepends a learnable class token and a histogram token before its attention layers §4 fig. 3.

Histogram prior (HPM). The motivation is that colour casts strongly affect underwater quality, and high-quality images tend to have an even colour distribution §1 §4.1. For each colour channel, a 64-bin histogram of the image is taken. The three histograms are flattened into a 192-length vector and passed through a linear layer to produce the histogram token §4.1 fig. 4. The paper does not say whether the counts are normalised §4.1.

Dynamic connection parallel block (DCPB). After the serial stages, parallel blocks exchange features between scales. The features of each scale first pass through a conv-attention layer on their own, then the other scales are resampled to that size by bilinear interpolation §4.2. Each scale's output is its own features plus two learned, weighted neighbours §4.2 eq. 1. A feed-forward layer follows §4.2. The paper contrasts this with direct, neighbour and dense connections §4.2 fig. 5.

Score. Each scale's class token goes through its own linear head. The final score is the average of these per-scale scores §4 fig. 3.

Learning to rank. Training draws a pair of images from the same group. The better image should receive a score at least a margin ε above the worse one. The margin is 0.5 §4.3 eq. 2. The paper does not define the loss for pairs with equal quality.

Enhancer supervision. The pretrained URanker gives a loss for an enhanced image: the sigmoid of its negated score §5.1 eq. 3. The enhancer's total loss is a pixel-wise content loss plus λ times this term §5.1 eq. 4. The content loss may be L1, L2 or perceptual §5.1. The paper does not say whether URanker is frozen during enhancer training.

Enhancer and normalization tail. The baseline enhancer, NU2Net, is a residual U-shape network built from convolution, instance normalization and ELU blocks §5.2 fig. 6. Its output goes through a normalization tail that rescales each channel's range to [0, 1] with min-max scaling, applied only where values overflow §5.2 eq. 5. The paper says this adds no parameters §6.4.2.

# URanker score for a batch  [§4] [fig. 3]
# Input:  image (B, 3, H, W), values in [0, 1]
# Output: score (B,), higher means better quality

def uranker_score(image):
    # Histogram prior: raw pixel counts, 64 bins per channel  [§4.1]
    # The image is zero-padded to a multiple of 32 first  [code: utils.py#L134-L139]
    histogram = channel_histograms(image, bins=64)   # (B, 192)
    maps = image
    stage_outputs = []
    for stage in stages:
        # Patch embedding downsamples; stage 1 uses patch size 4 in the code
        # [code: options/URanker.yaml#L14]  [§4]
        tokens = patch_embed(stage, maps)
        # Prepend a learnable class token and a histogram token  [§4] [§4.1]
        tokens = concat(class_token(stage), histogram_token(stage, histogram), tokens)
        # Conv-attention and feed-forward layers, applied alternately  [§4]
        tokens = serial_blocks(stage, tokens)
        stage_outputs.append(tokens)
        maps = to_map(drop_special_tokens(tokens))
    # Dynamic cross-scale fusion, repeated over the parallel blocks  [§4.2]
    # The released config uses 6 parallel blocks  [code: options/URanker.yaml#L17]
    stage_outputs = dynamic_fusion(stage_outputs)
    # Score from each scale's class token, then the mean over scales  [fig. 3]
    # The code scores stages 2 to 4 only  [code: all_model/URanker/model.py#L637]
    scores = [head(stage, class_token_of(stage_outputs[stage]))
              for stage in scored_stages]
    return mean(scores)

# One dynamic connection step over scales 2 to 4  [§4.2] [eq. 1]
# Scale 1 is not fused in the code  [code: all_model/URanker/model.py#L294]
# Scalars a1 to a6 are learned; the code starts each at 0.05
# [code: all_model/URanker/model.py#L229-L234]
def dynamic_fusion(f2, f3, f4):
    # Each scale adds two neighbours, resampled bilinearly to its size
    # [§4.2]  [code: all_model/URanker/model.py#L323-L326]
    g2 = f2 + a1 * resize(f3, to=f2) + a2 * resize(f4, to=f2)
    g3 = f3 + a3 * resize(f4, to=f3) + a4 * resize(f2, to=f3)
    g4 = f4 + a5 * resize(f3, to=f4) + a6 * resize(f2, to=f4)
    # The fused term joins the residual stream, then an MLP  [§4.2]
    return mlp(f2 + g2), mlp(f3 + g3), mlp(f4 + g4)

# Margin ranking loss for one pair  [§4.3] [eq. 2]
# better_score, worse_score: scalar outputs of the URanker
# The code clamps both scores to [-5, 5] first  [code: loss.py#L10-L11]
def pair_loss(better_score, worse_score):
    # Epsilon keeps the two scores apart; the paper sets it to 0.5  [§4.3]
    margin = 0.5
    return max(0, worse_score - better_score + margin)

# Enhancer training loss  [§5.1] [eq. 3] [eq. 4]
# enhanced, target: (B, 3, 256, 256)
def enhancer_loss(enhanced, target, uranker, weight):
    # Pixel loss; the released NU2Net config uses L1  [§5.1]
    # [code: options/NU2Net.yaml#L31]
    content = l1(enhanced, target)
    # Sigmoid of the negated URanker score, averaged over the batch  [eq. 3]
    # [code: loss.py#L16-L20]
    quality = mean(sigmoid(-uranker_score(enhanced)))
    # weight is lambda: 0.025 for NU2Net  [§6.3] [eq. 4]
    return content + weight * quality

# Normalization tail on the enhancer output  [§5.2] [eq. 5]
# The paper rescales a channel only where it overflows [0, 1]  [§5.2]
# The code rescales every image and channel when any value overflows
# [code: all_model/NU2Net/model.py#L4-L16]  [code: ...#L124-L125]
def normalization_tail(output):
    # Min-max over height and width, per image and channel  [eq. 5]
    low = min over height and width of output
    high = max over height and width of output
    # The code adds 1e-7 to the denominator; eq. 5 does not
    return (output - low) / (high - low + 1e-7)

Hyperparameters

Values from the paper are cited to the paper. Code values come from the released configs and are marked as such. Where the paper does not state a value, the row says so.

Results

Author-reported. Setup: URankerSet with the 800/90 group split, the same split for UIQA and UIE experiments §6.1. UIQA baselines use their released code; MUSIQ is retrained on the same data via IQA-PyTorch §6.2. Enhancers are trained on this split and evaluated with PSNR and SSIM §6.1 §6.3.

UIQA, SRCC and KRCC tab. 1.

The text reports gains over MUSIQ of 5.02% in SRCC and 6.84% in KRCC §6.2. These match the table (0.8655 / 0.8241 and 0.7402 / 0.6928).

Enhancers with and without the URanker loss, PSNR and SSIM tab. 2.

URanker design ablation, SRCC and KRCC tab. 3.

Normalization tail ablation, PSNR and SSIM tab. 4.

The text reports that the tail adds at least 2.456 dB PSNR and 0.039 SSIM over the other operators §6.4.2. These match the table against the best alternative (clip: 19.963 and 0.8837).

Paper vs code

The repository is RQ-Wu/UnderwaterRanker at commit f9bf3508121c5bd658a2f2c3fa8d4d3826ed1ce0. The paper's front matter links it, and its owner matches a listed author (Ruiqi Wu). The repository's LICENSE file is CC BY-NC 4.0, so the code is non-commercial. No code text is quoted here; differences are cited by path and line only. The shipped data are split lists, and the image directories and released weights are not in the repository URanker.yaml README.md.

• Enhancer epochs and schedule. The paper says 250 epochs with cosine annealing §6.1. The config sets 1000 epochs and CosineAnnealingLR with T_max 1000 NU2Net.yaml NU2Net.yaml. The scheduler steps after every batch runner.py. The resume logic assumes 50 batches per epoch runner.py. With 800 images and batch 16 that is 50 batches per epoch, so the cosine period is 1000 steps, about 20 epochs, and the rate rises again after each period (by reading the code; not run).
• Ranker epochs. The paper says 100 epochs §6.1. The config sets 1000 URanker.yaml. The enhancer config loads checkpoint_100best.pth NU2Net.yaml. The save rule writes file names from the epoch number, so this name covers epochs 50 to 99 runner.py.
• Ranker training input. The paper does not state an input size. The code halves each image side for ranker training and testing dataset.py dataset.py. In enhancer training, the ranker scores the enhancer's full-size 256 × 256 crops runner.py loss.py.
• Ranker batch size. Not stated in the paper. The code uses 1, and one epoch is 800 random pairs drawn from the 800 groups URanker.yaml dataset.py dataset.py.
• Pair loss. The paper's margin loss has no clamp §4.3 eq. 2. The code clamps both scores to [-5, 5] before the margin loss loss.py.
• Scales used for scoring. The paper says the final prediction averages the class-token scores of the different scales §4 fig. 3. The code scores stages 2 to 4 only, and stage 1 is not fused in the parallel blocks model.py model.py model.py.
• Histogram token. The paper describes one linear layer §4.1. The code uses four, one per stage model.py.
• Patch sizes. The paper says each serial block downsamples 2× §4. The code uses 4× in stage 1 and 2× in stages 2 to 4 URanker.yaml model.py.
• Scalar initialisation. The paper does not give an initial value for the DCPB scalars §4.2. The code starts all six at 0.05 model.py.
• Enhancer crop. The paper says inputs are cropped to 256 × 256 §6.1. The code uses a random resized crop dataset.py.
• Enhancer ground truth. The paper does not name the reference set for the enhancer's content loss §5.1 §6.3. The released config reads the UIEB raw-890 and reference-890 folders, with an 800/90 list split dataset.py uie_train_list.txt.
• Enhancer test input. The paper gives no test size. The code resizes each test image to a multiple of 16 dataset.py.
• Normalization tail. The paper says the tail acts only where values overflow, per channel §5.2 eq. 5. The code checks the whole batch for any overflow, then rescales every image and channel model.py model.py. The code adds 1e-7 to the denominator model.py. The enhancer's output has a Tanh and a residual add before the tail model.py model.py; the paper mentions only the residual layout, in the figure caption fig. 6.
• Enhancer model name. The NU2Net config asks for a ranker model named MyCoaT NU2Net.yaml. The loader imports all_model.<name>.model utils.py, and the only model folders are URanker and NU2Net [code: all_model/init.py#L1-L2]. Reading the code, the shipped enhancer config fails to import until the name is changed to URanker. We did not run it.
• Framework. The acknowledgements say a MindSpore implementation is provided § acknowledgements. The repository is PyTorch loss.py.

Pitfalls

• Test split used for checkpoint selection. Both training loops save the best epoch by the score on the 90-group test split: SRCC and KRCC for the ranker, PSNR and SSIM for the enhancer runner.py runner.py runner.py. There is no validation split in the code. The paper does not say how checkpoints were chosen §6.1, so the reported numbers may be optimistic. Re-run with a separate validation split before comparing against this paper.
• Scale mismatch in the enhancer loss. The ranker is trained on half-size images, but during enhancer training it scores full-size 256 × 256 crops dataset.py runner.py. The effect on the learned score is not measured here.
• Zero padding changes the histogram. The ranker zero-pads each image to a multiple of 32, centred, before the histogram is taken utils.py utils.py. Padding pixels are zero and fall into the lowest bin of each channel, so the first bin is inflated for images whose sides are not multiples of 32.
• Unnormalised histogram counts. The histogram uses raw counts over the padded image utils.py. The paper does not say whether the counts are normalised §4.1. The token then scales with image size, which the linear layer must absorb.
• Ties. Eq. 2 covers only q_n > q_m and q_n < q_m eq. 2. Equal quality is not defined. The bubble-sort labels avoid ties, but the paper does not state how equal-quality pairs were handled.
• Clamp in the pair loss. Scores outside [-5, 5] receive zero gradient through the clamp loss.py. The paper does not mention this clamp, and it can silence learning on saturated scores.
• Frozen ranker during enhancer training. The paper does not say whether URanker is frozen §5.1. In the code, its gradients are computed, but the enhancer's optimizer only holds the enhancer's parameters, so URanker is effectively fixed runner.py runner.py.
• Enhancer λ order. The paper gives λ = 0.025, 0.005 and 0.025 "respectively" for WaterNet, Ucolor and NU²Net §6.3. The sentence that introduces the three networks lists NU²Net first, so the mapping is ambiguous. Only the NU²Net value (0.025) is confirmed by the released config NU2Net.yaml.
• Which nine algorithms. The paper lists twelve algorithms, says three were dropped as near-duplicates, and does not say which three §3. The nine used are therefore not determined by the paper.
• Vote rule wording. The bubble-sort text describes the swap as moving the image that received the minority of votes §3. The intended rule is likely majority preference. With eleven volunteers, a pairwise majority cannot tie.
• Ablation percentage. The text says removing HPM reduces SRCC by 0.80% §6.4.1. Table 3 gives 0.8655 → 0.8585, which is 0.81% of the full model's SRCC. The difference is small but the text's value does not match the table.
• Dataset and weights. The image directories and released weights are not in the repository. The README refers to a checkpoint file that is not present README.md URanker.yaml. Reproducing the ranker needs the URankerSet images, which the card does not contain.

Lineage

• Builds on: the co-scale conv-attentional image Transformer (CoaT). This paper keeps its conv-attention in place of multi-head self-attention, for lower cost §4. The learning-to-rank setup follows RankIQA and RankSRGAN §4.3. Multi-scale Transformer IQA follows MUSIQ §1 §4.
• Compared against: UCIQE, UIQM, FDUM, NUIQ, Twice-Mixing and MUSIQ for UIQA §6.2; WaterNet and Ucolor as enhancers §6.3.
• Superseded by: not known from the content of this wiki.