Deep learning paper reproduction

Overview

Underwater photos lose colour and contrast because water absorbs and scatters light. Absorption varies with wavelength and often suppresses the red channel, leaving a blue or green cast, and scattering adds blur.

The paper we reproduced (Zhou et al., 2023) enhances underwater images with MFEF, a multi-view fusion network. Rather than learning from a single input, it gives the network several versions of the same image, each of which exposes a different kind of information, and lets it learn how to combine them. With no public code available, we reimplemented the method in PyTorch from the paper's description and trained it on the UIEB benchmark. This was a team of four for the TU Delft Deep Learning course. I built the model architecture, contributed to preprocessing and wrote the blog post.

Data and inputs

UIEB (Underwater Image Enhancement Benchmark) pairs 890 raw underwater images with reference images. All images are resized to 256 × 256 to give a fixed input size, and the pairs are split 85/15 into training and test sets.

Each raw image is turned into three views that enter the network as separate branches:

• Original. The unaltered image, which keeps all of the information the other two views may distort.
• White-balanced. A global gain derived from the scene mean rescales the image, which corrects the colour shift caused by water.
• CLAHE. Contrast-limited adaptive histogram equalisation is applied to the lightness channel in LAB colour space (8 × 8 tiles, clip limit 2.0), leaving the colour channels untouched. Equalising contrast tile by tile brings out local detail when visibility is low, and the clip limit stops noise in flat regions from being amplified.

Architecture

original, white-balanced, CLAHE views
  → residual enhancement module per view
  → multi-feature fusion
  → second residual stage, added back onto each view
  → shared convolutions at three scales
  → pixel-weighted channel attention per scale
  → coarse-to-fine reconstruction → enhanced image

The network is organised as an encoder of three parallel branches, a fusion stage, multi-scale attention and a decoder, with residual connections throughout.

Residual enhancement modules

Each view first passes through its own residual enhancement module (REM). A module is a stack of two residual blocks, each a sequence of four 3 × 3 convolutions that widen the representation to 64 and then 128 channels before projecting it back, with GELU activations and a skip connection that adds the block's input to its output. Because each block learns a correction to its input rather than a full mapping, the identity is easy to represent and gradients pass through deep stacks without degrading. The same module design is reused at every later stage of the network.

Multi-feature fusion

Three fusion modules combine the enhanced views. Within a module, each branch is convolved with a different kernel size (1 × 1, 3 × 3 and 5 × 5), so the three views are seen at different receptive fields, and each branch keeps a residual connection to its input. The branches are then concatenated along the channel axis and merged by a 3 × 3 convolution, so every fused map mixes information from all three views.

A second residual stage follows as a chain: the output of one REM is concatenated with the next fusion output before entering the next REM, so later modules see both fused features and the refinements made before them. Concatenation raises the channel count, while each result must be added pixel-wise to one of the three-channel input views. We therefore placed a 1 × 1 convolution after each of these REMs to project back to three channels, a choice the paper left open. Adding the result back onto the original, white-balanced and CLAHE images means that this stage learns a correction to each view rather than a new image.

Multi-scale channel attention

Each corrected view is then encoded at three scales. The same 3 × 3 convolution is applied three times with 2 × 2 max pooling in between, giving feature maps at full, half and quarter resolution. The paper's figure shows three convolution paths joined by unspecified dense connections. We read this as shared-kernel convolutions with downsampling and concatenated the same-scale maps from the three branches, which keeps the views aligned at each scale.

Each scale passes through a pixel-weighted channel attention module (PCAM), which reweights features in two ways:

• Pixel attention. A 3 × 3 convolution is gated element-wise by the sigmoid of a 1 × 1 convolution, so every spatial position receives its own weight.
• Channel attention. Global average pooling summarises each channel, and two fully connected layers with a sigmoid output turn that summary into a per-channel weight. Since channels come from different views, this decides how much each view contributes for a given image.

The module adds the pixel-weighted and channel-weighted maps to its input, so attention modulates the features rather than replacing them.

Reconstruction

The decoder works from the coarsest scale up. An REM refines the quarter-resolution features, which are upsampled by a factor of two, concatenated with the half-resolution attention output and refined again, and the same step is repeated at full resolution. A final 3 × 3 convolution aggregates the features into a three-channel image. This coarse-to-fine path combines the wider context captured at low resolution with fine detail from the full-resolution branch.

Loss and training

The loss combines an L1 pixel loss with a perceptual loss. For the perceptual term, the output and the reference both pass through an ImageNet-pretrained VGG-16, and the absolute differences between their activations are summed over its convolutional layers and added to the L1 term with a weight of 0.05. Comparing intermediate activations captures differences in texture, colour and structure that a pixel loss misses.

The model was optimised with Adam at a batch size of one for 200 epochs, as the paper proposes, on a single NVIDIA P100 GPU (about 12 hours). Loss on the held-out split was evaluated every 150 training images, and the weights with the lowest loss were kept. The model converged at epoch 89.

Evaluation

We evaluated with the paper's five metrics, which together cover fidelity to the reference and perceptual quality without one:

• PSNR measures pixel-level error against the reference.
• SSIM compares luminance, contrast and structure in local windows.
• PCQI is a patch-based contrast quality index built from mean intensity, signal strength and signal structure.
• UCIQE is a no-reference underwater metric combining chroma variation, luminance contrast and saturation.
• UIQM is a no-reference underwater metric combining measures of colourfulness (UICM), sharpness (UISM) and contrast (UIConM).

Each metric was computed per test image, and the distributions were examined as well as the means.

Results

The enhanced images come out visually close to the references. SSIM was concentrated in a narrow range, indicating that structure is preserved consistently across images, while PSNR varied more between images, with most of its mass at the higher values. PSNR and SSIM were lower than reported in the paper, and the other three metrics did not line up with the paper's values.

Key takeaways

• Reproducing a method without public code requires decisions on details that the paper leaves unspecified, and these decisions should be documented.
• The reproduction was visually close to the references but did not reach the scores reported in the paper.
• Per-image metric distributions were more informative than mean scores alone.