UIEC^2-Net: CNN-based Underwater Image Enhancement Using Two Color Space

In plain words

Underwater photos often have a bluish or greenish cast, low contrast and weak brightness. UIEC^2-Net is a small neural network that tries to fix this. Many earlier networks of this kind looked only at the usual red, green and blue values of each pixel. This one also looks at the same picture in HSV, a colour description with hue, saturation and value (brightness), which separates brightness and saturation from colour.

The network has three parts that are trained together. The first cleans up the picture in RGB and removes the colour cast. The second, working in HSV, adjusts brightness and saturation across the whole image using simple curves that it learns. The third decides, pixel by pixel, how much to trust each of the first two results.

The idea is that the two colour spaces cover each other's weak spots. RGB is good at removing colour cast but does not expose brightness and saturation directly. HSV does, but its hue channel is sensitive and can distort colours in some areas. The per-pixel mix lets the network take the good parts of each.

TL;DR

• A three-block CNN for underwater image enhancement: an RGB block, an HSV block that predicts global curves, and a per-pixel attention block that blends the two §III-A.
• Trained end-to-end with L1, SSIM, an HSV loss and a VGG-19 perceptual loss §III-B.
• Author-reported: best full-reference scores (MSE, PSNR, SSIM) on the synthetic and real-world test sets, and best UCIQE, UISM and UIQM on real-world images tab. I tab. II tab. III.
• The released PyTorch code trains with SSIM alone and differs from the paper in several layers and losses. Read ## Paper vs code before implementing.

Problem

Deep underwater enhancers built only on RGB do not restore brightness and saturation well, and the paper argues that RGB cannot directly expose those properties §I. HSV exposes them, but the paper notes that its hue channel is sensitive and can distort colours in some areas §I §III-A; the paper's answer is to use both spaces. Underwater images are often bluish or greenish, low in brightness, yellowish or hazy §IV-C. Physical-model methods assume a simplified formation model that fails across scenes §I. Paired training data are scarce: no real underwater image has a ground-truth reference, so most training uses synthetic images §II §IV-B.

Method

The paper uses IEEE numbering, so sections are roman numerals with lettered subsections (for example §III-A), tables are roman (tab. I to IV) and figures and equations are arabic.

1. RGB pixel-level block §III-A. A plain CNN with no down-sampling. The paper describes eight 3×3 convolutions with stride 1. Each uses batch normalisation, and the first seven are followed by LeakyReLU. The last is followed by a sigmoid so the output lies in [0, 1]. Output: an RGB image of the same size.
2. RGB to HSV §III-C eq. 8 eq. 9 eq. 10. The RGB block output is converted to HSV with formulas built from max and min operations, so gradients can flow through them.
3. HSV global-adjust block §III-A. Five 3×3 convolutions with stride 1, each followed by LeakyReLU. A 2×2 max-pool follows the first four and not the fifth. Global average pooling, then a fully connected layer, regresses the knot values of four piecewise-linear curves.
4. Four curves §III-A eq. 1. The curves are: value on value (luminance), saturation on saturation, saturation on hue, and hue on hue (colour). Each applies the piecewise-linear mapping in eq. 1 to its channel. The curve shape is set by its knot values.
5. HSV to RGB §III-C eq. 11 eq. 12 eq. 13 eq. 14. The adjusted HSV image is converted back to RGB. The paper replaces the usual formulas with piecewise-linear ones that are differentiable.
6. Attention block §III-A. It takes the raw input image, the RGB block output and the HSV block output concatenated along channels (9 channels). The paper describes eight 3×3 convolutions with batch normalisation and LeakyReLU, and a sigmoid on the last. It outputs 6 maps in [0, 1]: the first three weight the RGB output, the last three weight the HSV output.
7. Combine §III-A. Each weight map is multiplied with its output and the two results are added. The sum is the final enhanced image.

Key equations, as printed. The curve, with δ defined in §III-A as 0 for x < 0, x for 0 ≤ x ≤ 1 and 1 for x > 1:

S(\hat I_i^{jl}) = k_0 + \sum_{m=0}^{M-1} (k_{m+1} - k_m)\,\delta(M\hat I_i^{jl} - m) eq. 1

The total loss, with the weights in §III-B:

\mathcal{L} = (\lambda_{pixel} + \lambda_{whole})(\omega_{\ell 1}\mathcal{L}_{\ell 1} + \omega_{SSIM}\mathcal{L}_{SSIM}) + \omega_{hsv}\mathcal{L}_{hsv} + \omega_{perc}\mathcal{L}^{\phi}_{j} eq. 2

The HSV loss, computed in the conical HSV space, with hue in [0, 2π):

\mathcal{L}_{hsv} = \left\| \hat S \hat V \cos(\hat H) - S V \cos(H) \right\|_1 eq. 3

Pseudocode for the forward pass. The paper has no algorithm box, so this is a reading of §III-A, §III-B and §III-C:

# UIEC^2-Net forward pass, from the text of [§III-A] and [§III-C]
# Input:  image (B, 3, H, W) with values in [0, 1]
# Output: enhanced image (B, 3, H, W)

def uiec2_forward(image):
    # 1. RGB pixel-level block: plain CNN, no down-sampling  [§III-A]
    rgb_out = rgb_block(image)                    # (B, 3, H, W), in [0, 1]

    # 2. Convert the RGB block output to HSV  [§III-C] [eq. 8] [eq. 9] [eq. 10]
    hsv_in = rgb_to_hsv(rgb_out)                  # (B, 3, H, W)

    # 3. HSV global-adjust block: convolutions, global average pooling,
    #    then a linear layer that regresses the knots of four curves  [§III-A]
    knots = hsv_block(hsv_in)                     # (B, 4, M)

    # 4. Apply the four curves to value, saturation (two curves) and hue
    #    [§III-A] [eq. 1]
    hsv_adj = apply_curves(hsv_in, knots)         # (B, 3, H, W)

    # 5. Convert the adjusted HSV back to RGB with piecewise-linear
    #    functions in place of the usual formulas  [§III-C] [eq. 11]
    hsv_rgb = hsv_to_rgb(hsv_adj)                 # (B, 3, H, W)

    # 6. Attention block sees the raw image and both outputs (9 channels)
    #    and predicts 6 maps in [0, 1]: 3 for RGB, 3 for HSV  [§III-A]
    attention = attention_block(
        concat([image, rgb_out, hsv_rgb], dim=1)) # (B, 6, H, W)

    # 7. Weight each output per pixel and per channel, then add  [§III-A]
    rgb_weight = attention[:, 0:3]                # (B, 3, H, W)
    hsv_weight = attention[:, 3:6]                # (B, 3, H, W)
    return rgb_weight * rgb_out + hsv_weight * hsv_rgb


# Piecewise-linear curve with M knots  [§III-A] [eq. 1]
# x: one channel of an image, values in [0, 1]
# knots: M knot values k_0 ... k_(M-1) for this image and channel
def apply_curve(x, knots):
    # Map each value onto the knot axis, 0 .. M  [eq. 1]
    position = M * x
    # Start at k_0, then add each knot step ramped in by clamp to [0, 1]
    # over its interval  [eq. 1]
    out = knots[0]
    for m in range(M - 1):
        step = knots[m + 1] - knots[m]
        out = out + step * clamp(position - m, 0, 1)
    return out

Pseudocode for the losses, with the weights from §III-B. Eq. 2 is shown as printed; §III-B gives the values:

# Total loss  [§III-B] [eq. 2]
# Weights: w_l1 = 1, w_ssim = 1, w_hsv = 1, w_perc = 0.5  [§III-B]
def total_loss(prediction, gt):
    # lambda_pixel + lambda_whole: 0.5 + 0.5 for epochs 1-20, 0.1 + 0.9 after
    # [§III-B]; both pairs sum to 1
    lambda_sum = 1.0
    # L1 and SSIM terms on RGB pixels  [eq. 4] [eq. 6]
    rgb_terms = 1.0 * l1(prediction, gt) + 1.0 * ssim_loss(prediction, gt)
    # HSV loss in the conical HSV space  [eq. 3]
    hsv_term = 1.0 * hsv_loss(prediction, gt)
    # Perceptual loss on VGG-19 layer 4_3  [eq. 7]
    perc_term = 0.5 * perceptual_loss(prediction, gt)
    return lambda_sum * rgb_terms + hsv_term + perc_term  # [eq. 2]


# SSIM loss over 11x11 patches  [§III-B] [eq. 5] [eq. 6]
def ssim_loss(prediction, gt):
    # The paper computes SSIM on gray images; the conversion is not stated
    gray_p = to_gray(prediction)
    gray_g = to_gray(gt)
    # Per-pixel means, variances and covariance over the patch  [eq. 5]
    # c1 = 0.02 and c2 = 0.03 as printed  [§III-B]
    ssim_map = ssim_per_pixel(gray_p, gray_g, c1=0.02, c2=0.03)
    return 1 - mean(ssim_map)                     # [eq. 6]


# HSV loss  [§III-B] [eq. 3]
def hsv_loss(prediction, gt):
    # Hue in radians [0, 2pi); saturation and value in [0, 1]  [§III-B]
    hue_p, sat_p, val_p = rgb_to_hsv_radians(prediction)
    hue_g, sat_g, val_g = rgb_to_hsv_radians(gt)
    # Per-pixel conical term S * V * cos(H)  [eq. 3]
    cone_p = sat_p * val_p * cos(hue_p)
    cone_g = sat_g * val_g * cos(hue_g)
    # L1 norm of the difference; reduction is not stated  [eq. 3]
    return l1_norm(cone_p - cone_g)


# Perceptual loss on VGG-19 features  [§III-B] [eq. 7]
def perceptual_loss(prediction, gt):
    # Layer 4_3 of an ImageNet-pretrained VGG-19  [§III-B]
    feat_p = vgg19_layer_4_3(prediction)          # (B, C, H', W')
    feat_g = vgg19_layer_4_3(gt)                  # (B, C, H', W')
    # Per-image norm of the feature difference, summed over the batch and
    # divided by C * H * W; the norm type is not stated  [eq. 7]
    return sum_over_batch(norm(feat_p - feat_g)) / (C * H * W)

Colour conversions, in words. RGB to HSV uses V = max(R, G, B), S = (V − min) / V with S = 0 when V = 0, and a hue that is 60 times a ratio on whichever channel is the maximum, with 120 and 240 offsets, wrapped into [0, 360] eq. 8 eq. 9 eq. 10. HSV to RGB uses the value, saturation and hue with clamped linear terms δ over [0, 60] and a 360-degree hue scale §III-C eq. 11 eq. 12 eq. 13 eq. 14. The printed set has a typo: eq. 12 and eq. 13 are the same B formula, and no G formula is printed. See ## Pitfalls.

Hyperparameters

Values marked "not stated" are absent from the paper. Code values are from config/UIEC2Net.py and the modules it names. The code is a different configuration from the paper in several places, see ## Paper vs code.

Results

Author-reported. Setup: 2000 training images and a test set of 90 real and 900 synthetic images; baselines run with the authors' pretrained weights, not retrained §IV-A §IV. The paper marks the best value in red and the second best in blue. Each row below is one method; "Source" is the table in the paper.

Synthetic images, full-reference metrics tab. I:

Real-world images, full-reference metrics tab. II:

Real-world images, no-reference metrics tab. III:

Ablation on real-world images tab. IV:

Headline claims, as the paper states them: the full model is best on the three full-reference metrics on both sets §IV-B §IV-C; best on UCIQE, UISM and UIQM among the no-reference metrics §IV-C; and the HSV and attention blocks cost some UCIQE and UIQM but improve the full-reference scores §IV-D.

Paper vs code

Code: https://github.com/BIGWangYuDong/UWEnhancement at e012263575cb4ea22db3ec40d28825797745d067, linked from the abstract. The README names the paper and lists the corresponding author's contact address, so the repository is the authors' own. The repository has no LICENSE file, so it can be cited by path and line only.

• Paper says the PyTorch code is a conversion of the authors' TensorFlow code, and that the released checkpoints were transcoded from TensorFlow README.md. The paper does not mention TensorFlow.
• RGB block: paper has eight 3×3 convolutions with batch normalisation §III-A. Code has six 3×3 convolutions, then one 1×1 convolution, with InstanceNorm and ReLU on the last two 3×3 layers UIEC2Net.py UIEC2Net.py UIEC2Net.py.
• HSV block input: paper converts the RGB block output to HSV §III-C. Code concatenates that HSV tensor with itself, giving 6 input channels UIEC2Net.py UIEC2Net.py.
• HSV block pooling: paper pools after the first four convolutions §III-A. Code pools after the first three and not after the fourth or fifth UIEC2Net.py.
• Saturation curve: paper names two saturation curves (on saturation and on hue) without saying how they combine §III-A. Code averages them UIEC2Net.py.
• Output clamp: code clamps the saturation and value curves to [0, 1] but not hue UIEC2Net.py.
• Attention block: paper uses eight 3×3 convolutions §III-A. Code has six 3×3 and one 1×1, with a 9-channel input in the order image, RGB output, HSV output UIEC2Net.py UIEC2Net.py.
• Output combination: paper multiplies each weight map with its output and adds the two §III-A. Code also multiplies each term by 0.5 UIEC2Net.py.
• HSV loss: paper uses cos on both the prediction and the ground truth eq. 3. Code uses sin on the ground truth hsv_loss.py. Code also scales hue by 2π and stores hue in [0, 1] hsv_loss.py RGB2HSV.py.
• HSV to RGB: code computes G as well, from the same clamped-linear pattern as R and B HSV2RGB.py HSV2RGB.py. The paper prints no G formula §III-C.
• SSIM: paper says gray images, 11 × 11 patches, c1 = 0.02 and c2 = 0.03 §III-B. Code runs SSIM on each RGB channel with a Gaussian window of size 11 and σ = 1.5, and uses C1 = 0.01² and C2 = 0.03² ssim_loss.py ssim_loss.py.
• Loss used for training: paper's Eq. 2 combines L1, SSIM, HSV and perceptual terms eq. 2. The shipped training script computes only the SSIM loss on the network output train.py. Its default --config is config/UWCNN.py train.py, so config/UIEC2Net.py must be passed explicitly.
• Loss weights in config: L1 weight 2.0 and perceptual weight 0, against the paper's 1 and 0.5 UIEC2Net.py. The config has no λ_pixel or λ_whole.
• Training data and crop: paper resizes to 350 and crops to 320 §IV-A. The config's training pipeline has only load, random flip and tensor conversion UIEC2Net.py.
• Learning rate and epochs: paper 1e-4 and 50 epochs §IV-A. Config 1e-3 and 1000 epochs UIEC2Net.py UIEC2Net.py.
• Test-time augmentation: the repository has an eight-way augmentation script test_TTAx8.py. The paper does not say whether Tables I to IV used it.
• Device: the model and the curve functions call .cuda() directly UIEC2Net.py UIEC2Net.py UIEC2Net.py. Without CUDA the code should fail as written (inferred from the code, not run).
• Weights: the clone has no checkpoint file. The test script loads checkpoints/UIEC2Net.pth test_UIEC2Net.py, and checkpoints/README.md only says "put checkpoints here". The README points to external download links, which were not checked.
• Data paths: the config points at absolute paths on the authors' machine UIEC2Net.py.

Pitfalls

• Curve sum limit. Eq. 1 sums m from 0 to M−1 and uses k_{m+1}, so the last term needs k_M, which is not defined. The code sums M−1 steps, from 0 to M−2 UIEC2Net.py. The pseudocode follows the code. M is not stated in the paper; the code uses 11 UIEC2Net.py.
• Loss split. §III-B says the L1 and SSIM terms are applied once to the RGB block output and once to the final output, with weights λ_pixel and λ_whole. Eq. 2 shows neither split: it multiplies both terms by the sum λ_pixel + λ_whole. Both schedules sum to 1, so the factor in Eq. 2 is 1 at every epoch. An implementer has to pick the split; the paper does not say which output the HSV and perceptual terms use.
• Attention channel order. The sentence that combines the outputs lists the HSV output first, but the channel description says the first three channels weight the RGB output §III-A. The pseudocode follows the channel description, as the code does UIEC2Net.py. The 0.5 factor in the code is not in the paper.
• HSV to RGB. eq. 12 and eq. 13 are the same B formula, and no G formula is printed §III-C. The code's G formula follows the R and B pattern. Mark it as code-derived.
• HSV loss sign. The paper uses cos on both terms eq. 3. The code uses sin on the ground truth hsv_loss.py. The HSV loss is not in the shipped training path, so the code does not test it.
• Hue units. The paper gives hue in degrees in eq. 10 and in radians in [0, 2π) for the loss §III-B. The code stores hue as a fraction of 360. Keep one convention throughout.
• Division and epsilon. Eq. 9 and eq. 10 divide by V and by V − min, and the paper says nothing for the grey case. The code adds 1e-4 to both denominators RGB2HSV.py RGB2HSV.py.
• SSIM constants and colour. c1 = 0.02 is printed in the paper. The code uses C1 = 0.01². Which value is intended is not stated. The paper says gray images but gives no conversion; the code works per RGB channel.
• Data pipeline. The 350 resize and 320 crop from §IV-A are missing from the config, so a faithful run needs a custom pipeline. The synthetic generator follows UWCNN and is not in this repository §IV-A.
• Baselines. Paper uses only the type-1 UWCNN model, though UWCNN has ten type-specific models §IV-B. The baselines are the authors' pretrained weights and were not retrained on the same data §IV. The UWCNN row in tab. III almost equals the raw-input row (2.55656 against 2.5656 for UICM), which may be a typo. Treat that row with care.
• UIQM formula. Section IV-C writes UIQM = c1 × UICM + c2 × UISM + c3 × UIQM, which refers to itself. The third term must be UIConM, as the text describes it.
• Table IV text. The text says the HSV and attention blocks decrease UCIQE and UIQM, but the RGB-only variant has the highest UCIQE and the variant without attention has the highest UIQM tab. IV.
• Cost is not stated. The paper says the model has far fewer parameters than deep CNNs §I but gives no count, no runtime and no memory figure. The code's get_parameters prints a count at run time; the number is not in the paper.
• Reproducing numbers. The clone has no weights and no evaluation outputs, and the test data are not in the repository. The paper's numbers cannot be checked against the code here.

Lineage

• Builds on: DIFAR (Moran and Slabaugh), arXiv:1911.13175, for the neural curve layer and the conical HSV loss. The paper simplifies DIFAR's model §I §III-B.
• Related global-transform idea: DeepISP (Schwartz et al.), named in §I. No arXiv id in the source bibliography.
• Compared against: HE, WB, UDCP, ULAP, UWGAN, UGAN, UWCNN and DUIENet §IV-B §IV-C.
• Dataset: the UIEBD benchmark of Li et al. (IEEE TIP 2019), which the paper uses for real-world training and testing §IV-A. No arXiv id in the source bibliography.
• Superseded by: not known from content already in this wiki.