Underwater Image Enhancement with Cascaded Contrastive Learning

TL;DR

The proposed CCL-Net involves two cascaded stages, i.e., a color correction stage tailored to the color deviation issue and a haze removal stage tailored to improve the visibility and contrast of underwater images, and applies contrastive loss as an additional constraint to guide the training of each stage.

Abstract

Underwater image enhancement (UIE) is a highly challenging task due to the complexity of underwater environment and the diversity of underwater image degradation. Due to the application of deep learning, current UIE methods have made significant progress. Most of the existing deep learning-based UIE methods follow a single-stage network which cannot effectively address the diverse degradations simultaneously. In this paper, we propose to address this issue by designing a two-stage deep learning framework and taking advantage of cascaded contrastive learning to guide the network training of each stage. The proposed method is called CCL-Net in short. Specifically, the proposed CCL-Net involves two cascaded stages, i.e., a color correction stage tailored to the color deviation issue and a haze removal stage tailored to improve the visibility and contrast of underwater images. To guarantee the underwater image can be progressively enhanced, we also apply contrastive loss as an additional constraint to guide the training of each stage. In the first stage, the raw underwater images are used as negative samples for building the first contrastive loss, ensuring the enhanced results of the first color correction stage are better than the original inputs. While in the second stage, the enhanced results rather than the raw underwater images of the first color correction stage are used as the negative samples for building the second contrastive loss, thus ensuring the final enhanced results of the second haze removal stage are better than the intermediate color corrected results. Extensive experiments on multiple benchmark datasets demonstrate that our CCL-Net can achieve superior performance compared to many state-of-the-art methods. The source code of CCL-Net will be released at https://github.com/lewis081/CCL-Net.

Figures (13)

• Figure 1: The enhanced results by proposed CCL-Net. (a) is the raw underwater images with a predominant greenish tone. (b) is the pseudo-reference counterpart. (c) is the output of the CC-Net. (d) is the output of the HR-Net, i.e., the final result of CCL-Net.
• Figure 2: Architecture of the proposed CCL-Net. It consists a color correction stage and a haze removal stage. In the first color correction stage, a raw underwater image $\mathbf{I}$ is first transformed into the Lab color space, resulting in a luminance channel $\mathbf{I_L}$ and two color channels $\mathbf{I_{ab}}$. Then, the luminance channel $\mathbf{I_L}$ is discarded before feeding into the CC-Net which focuses on color correction only. As a result, only the two color channels $\mathbf{I_{ab}}$ are fed into the CC-Net, resulting in corrected color channels $\mathbf{I_{ab}^{CC}}$. Finally, the corrected color channels $\mathbf{I_{ab}^{CC}}$ and the original luminance channel $\mathbf{I_L}$ are merged and converted back into a color corrected yet hazy underwater image in the RGB format $\mathbf{I^{CC}}$. A hybrid loss consisting of an explicit color correction loss and a contrastive loss is used as the constraint of CC-Net. In the second haze removal stage, the color corrected yet hazy underwater image from the previous stage $\mathbf{I^{CC}}$ is fed into the HR-Net, resulting in an enhanced result with clear structures and fine details. Similarly, a hybrid loss consisting of an explicit structural similarity (SSIM) loss and a contrastive loss is used as the constraint of HR-Net. This helps the enhanced image approach the reference image and get away from the color corrected yet hazy image, ultimately achieving a visually appealing enhanced image with high color saturation and contrast.
• Figure 3: The schematic illustration of CC-Net. The CC-Net learns the difference between raw color channels and reference color channels to achieve color correction.
• Figure 4: The schematic illustration of HR-Net. An attention-based multi-scale information fusion module is utilized in HR-Net to select appropriate scale features for regions with varying haze concentration.
• Figure 5: Visual comparisons on underwater images with pronounced greenish and bluish tone from UIEB-T90. The numbers on the top-left corner of each image refer to its PSNR/SSIM values.