Underwater Image Enhancement via Medium Transmission-Guided Multi-Color Space Embedding

Compared with terrestrial scene images, the color deviations of underwater images cover more comprehensive ranges, differing from the bluish or greenish tone to a yellowish one. The diversity in color casts severely limits the traditional network architectures [41, 42]. Inspired by the traditional enhancement algorithms that operate in various color spaces [23, 43], we extract features in three color spaces (RGB, HSV, and Lab) where the same image has different visual representations in various color systems as shown in Fig. 3.

Concretely, the image is easy to store and display in RGB color space because of its strong color physical meaning. However, the three components (R, G, and B) are highly correlated, which are easy to be affected by the changes of luminance, occlusion, shadow, and other factors. By contrast, HSV color space can intuitively reflect the hue, saturation, brightness, and contrast of the image. Lab color space makes the colors better distributed, which is able to express all the colors that the human eye can perceive.

These color spaces have obvious differences and advantages. To combine their properties for underwater image enhancement, we incorporate the characteristics of different color spaces into a unified deep structure, where all the image degradation related components (color, hue, saturation, intensity, and luminance) can be taken into account. Moreover, the color difference of two points with a small distance in one color space may be large in other color spaces. Thus, the multiple color spaces embedding can facilitate the measurement of color deviations of underwater images. Additionally, the multi-color space encoder brings more nonlinear operations during color space transformation. It is known that the nonlinear transformation generally improves the performance of deep models [44]. In the ablation study, we will analyze the contribution of each color space.

Fig. 4 presents the details of the residual-enhancement module. This residual-enhancement module aims to preserve the data fidelity and address gradient vanishing [45]. In each residual-enhancement module, the convolutional layers have an identical number of filters. The numbers of filters are progressively increased from 128 to 512 by a factor 2 in the encoder network while they are decreased from 512 to 128 by a factor 2 in the decoder network. All the convolutional layers have the same kernel sizes of 3$\times$3 and stride 1.

In view of the specific definition of each color space, these features extracted from three color spaces should have different contributions. Therefore, we employ a channel-attention module to explicitly exploit the interdependencies between the channel features extracted from different color spaces. The details of the channel-attention module are depicted in Fig. 5.

Assume the input features $\mathcal{F}=\textsf{Cat}(\mathbf{F}_{1},\mathbf{F}_{2},\cdots,\mathbf{F}_{N})$ $\in$ $\mathbb{R}^{N\times H\times W}$, where $\mathbf{F}$ is a feature map from one path at a specific level (the level is denoted as 1, 2, and 3 in Fig. 2), $N$ is the number of feature maps, Cat represents the feature concatenation; and $H$ and $W$ are the height and width of input image, respectively. We first perform the global average pooling on input features $\mathcal{F}$, leading to a channel descriptor $\mathbf{z}\in\mathbb{R}^{N\times 1}$, which is an embedded global distribution of channel-wise feature responses. The $k$-th entry of $\mathbf{z}$ can be expressed as:

where $k\in[1,N]$. To fully capture channel-wise dependencies, a self-gating mechanism [46] is used to produce a collection of per-channel modulation weights $\mathbf{s}\in\mathbb{R}^{N\times 1}$:

where $\mathscr{\sigma}(\cdot)$ represents the Sigmoid activation function, $\mathscr{\delta}(\cdot)$ represents the ReLU activation function, $\ast$ denotes the convolution operation, and $\mathbf{W}_{1}$ and $\mathbf{W}_{2}$ are the weights of two fully-connected layers with the numbers of their output channels equal to $\frac{N}{r}$ and $N$, respectively, where $r$ is set to 16 for reducing the computational costs. At last, these weights are applied to input features $\mathcal{F}$ to generate rescaled features $\mathcal{U}$ $\in$ $\mathbb{R}^{N\times H\times W}$. Moreover, to avoid gradient vanishing problem and keep good properties of original features, we treat the channel-attention weights in an identical mapping fashion:

where $\oplus$ and $\otimes$ denote the pixel-wise addition and pixel-wise multiplication, respectively.

According to the image formation model in bad weather [47, 48], which is widely used in image dehazing and underwater image restoration algorithms [2, 29, 30], the quality-degraded image can be expressed as:

where $x$ indicates the pixel index, $I$ is the observed image, $J$ is the clear image, $A$ is the homogeneous background light, and $T$ is the medium transmission that represents the percentage of scene radiance reaching the camera after reflecting in the medium, indicating the degrees of quality degradation in different regions.

We incorporate the medium transmission map into the decoder network via the proposed medium transmission guidance module. Specifically, we use the reverse medium transmission (RMT) map (denoted as $\overline{T}\in\mathbb{R}^{H\times W}$) as the pixel-wise attention map. The RMT map $\overline{T}$ is obtained by $\mathbf{1}$-$T$ ($T\in\mathbb{R}^{H\times W}$ is the medium transmission map in the range of [0,1], and $\mathbf{1}\in\mathbb{R}^{H\times W}$ is the matrix with all elements equal to 1), which indicates that the higher quality degradation pixels should be assigned larger attention weights.

Since the corresponding ground truth medium transmission map of an input underwater image is not available in practice, it is difficult to train a deep neural network for the estimation of medium transmission map. To solve this issue, we employ prior-based estimation algorithms to obtain the medium transmission map. Inspired by the robust general dark channel prior [30], we estimate the medium transmission map as:

where $\tilde{T}$ is the estimated medium transmission map, $\Omega(x)$ represents a local patch of size 15$\times$15 centered at $x$, and $c$ denotes the color channel. As shown, the medium transmission estimation is related to the homogeneous background light $A$. In [30], the homogeneous background light is estimated based on the depth-dependent color change. Due to the limited space, we refer the readers to [30] for more details. We will compare and analyze the effects of the medium transmission maps estimated by different algorithms in the ablation study.

With the RMT map, the schematic illustration of the proposed medium transmission guidance module is shown in Fig. 6. As shown, we utilize the RMT map as a feature selector to weight the importance of different spatial positions of the features. The high-quality degradation pixels (the pixels with larger RMT values) are assigned higher weights, which can be expressed as:

where $\mathcal{V}$ $\in$ $\mathbb{R}^{M\times H\times W}$ and $\mathcal{U}$ $\in$ $\mathbb{R}^{M\times H\times W}$ respectively represent the output features after the medium transmission guidance module and the input feature. We treat the RMT weights as an identity connection to avoid gradient vanishing and tolerate the errors caused by inaccurate medium transmission estimation. Besides, our purely data-driven framework also tolerates the inaccuracy of medium transmission maps.

Following previous works [49, 50], to achieve a good balance between visual quality and quantitative scores, we use the linear combination of the $\ell_{2}$ loss $L_{\ell_{2}}$ and the perceptual loss $L_{per}$, and the final loss $L_{f}$ for training our network is experessed as:

where $\lambda$ is empirically set to 0.01 for balancing the scales of different losses. Specifically, the $\ell_{2}$ loss measures the difference between the reconstructed result $\hat{J}$ and corresponding ground truth $J$ as:

The perceptual loss is computed based on the VGG-19 network $\phi$ [44] pre-trained on the ImageNet dataset [51]. Let $\phi_{j}(\cdot)$ be the $j$th convolutional layer. We measure the distance between the feature representations of the reconstructed result $\hat{J}$ and ground truth image $J$ as:

We compute the perceptual loss at layer relu5_4 of the VGG-19 network. An ablation study towards the loss function will be presented.

To train the Ucolor, we randomly selected 800 pairs of underwater images from UIEB [18] underwater image enhancement dataset. The UIEB dataset includes 890 real underwater images with corresponding reference images. Each reference image was selected by 50 volunteers from 12 enhanced results. It covers diverse underwater scenes, different characteristics of quality degradation, and a broad range of image content, but the number of underwater images is inadequate to train our network. Thus, we incorporated 1,250 synthetic underwater images selected from a synthesized underwater image dataset [12], which includes 10 subsets denoted by ten types of water (I, IA, IB, II, and III for open ocean water and 1, 3, 5, 7, and 9 for coastal water). To augment the training data, we randomly cropped image patches of size 128$\times$128.

We implemented the Ucolor using the MindSpore Lite tool [52]. A batch-mode learning method with a batch size of 16 was applied. The filter weights of each layer were initialized by Gaussian distribution, and the bias was initially set as a constant. We used ADAM for network optimization and fixed the learning rate to $1e^{-4}$.

Since the source code of Ancuti et al. [9] is not publicly available, we used the code²²2https://github.com/bilityniu/underimage-fusion-enhancement implemented by other researchers to realize code of [9]. For Li et al. [4], Peng et al.  [29], GDCP [30], UcycleGAN [10], and Water-Net [18], we used the released codes to produce their results. The results of Guo et al. [17] were provided by the authors. Note that UcycleGAN [10] is an unsupervised methods, i.e., training with unpaired data, and thus there is no need to retrain it with our training data. Same as our Ucolor, Water-Net [18] randomly selected the same number of training data from the UIEB dataset for training. For UWCNN [12], we used the original UWCNN models, in which each UWCNN model was trained using the underwater images synthesized by one type of water. We discarded the UWCNN$\_$typeIA and UWCNN$\_$typeIB models because their results are similar to those of UWCNN$\_$typeI. Besides, we also retrained the UWCNN model (denoted as UWCNN$\_$retrain) using the same training data as our Ucolor.

For Test-C60 and SQUID that do not have corresponding ground truth images, we employ the no-reference evaluation metrics UCIQE [54] and UIQM [55] to measure the performance of different methods. A higher UCIQE or UIQM score suggests a better human visual perception. Please note that the scores of UCIQE and UIQM cannot accurately reflect the performance of underwater image enhancement methods in some cases. We refer the readers to [18] and [53] for the discussions and visual examples. In our study, we only provide the scores of UCIQE and UIQM as the reference for the following research. In addition, we also provide the scores of NIQE [56] of different methods as the reference though it was not originally devised for underwater images. A lower NIQE score suggests better image quality. Although SQUID provides a script to evaluate color reproduction and transmission map estimation for the underwater image, this script requires the estimated transmission map or the results having a resolution of 1827$\times$2737. However, the compared deep learning-based methods do not need to estimate the transmission maps. Moreover, our current GPU cannot process the input image with such a high resolution. Besides, the sizes of the color checker in the images of SQUID are too small to be cropped for full-reference evaluations. Instead, we resized the images in the SQUID testing dataset to a size of 512$\times$512 and processed them by different methods. We conducted a user study to measure the perceptual quality of results on Test-C60 and SQUID. Specifically, we invited 20 human subjects to score the perceptual quality of the enhanced images independently. The scores of perceptual quality range from 1 to 5 (worst to best quality). These subjects were trained by observing the results from 1) whether the results introduce color deviations; 2) whether the results contain artifacts; 3) whether the results look natural; and 4) whether the results have good contrast and visibility.

For Color-Check7, we measured the dissimilarity of color between the ground-truth Macbeth Color Checker and the corresponding enhanced results. To be specific, we extracted the color of 24 color patches from the ground-truth Macbeth Color Checker. Then, we respectively cropped the 24 color patches for each enhanced result and computed the average color values of each color patch. At last, we followed the previous method [25] to employ CIEDE2000 [57] to measure the relative perceptual differences between the corresponding color patches of ground-truth Macbeth Color Checker and those of enhanced results. The smaller the CIEDE2000 value, the better.

In this section, we conduct visual comparisons on diverse testing datasets. We refer readers to the Google Drive link³³3https://drive.google.com/file/d/1zrynw05ZgkVMAybGo9Nhqg2mmIYIMVOT/view?usp=sharing for all results of different methods (about 7.4 GB).

We first show the comparisons on a synthetic image in Fig. 7. The competing methods either fail to dehaze the input image or they introduce undesirable color artifacts. All the methods under comparison fail to recover the complete scene structure. Our result (Fig. 7(l)) is closest to the ground-truth images and obtains the best PSNR/MSE scores. Furthermore, the RMT map shown in Fig. 7(b) indicates that the high degradation regions have large pixel values. With the RMT map, our method highlights these regions and enhances them well.

We then show the results of different methods on a real underwater image with obvious greenish color deviation in Fig. 8. In Fig. 8(a), the greenish color deviation significantly hides the structural details of the underwater scene. In terms of color, GDCP [30], UcycleGAN [10], UWCNN_typeII [12] and UWCNN_retrain [12] introduce extra color artifacts. All the compared methods under-enhance the image or introduce the over-saturation. In comparison, our Ucolor effectively removes the greenish tone and improves the contrast without obvious over-enhancement and over-saturation.Although UWCNN_retrain [12], Unet-U [41], and Unet-RMT [41] were trained with the same data as our Ucolor, their performance is not as good as our Ucolor, which demonstrates the advantage of our specially designed network structure for underwater image enhancement.

We also show comparisons on challenging underwater images sampled from Test-C60 in Fig. 9. These underwater images suffer from high backscattering and color deviations as shown in Fig. 9(a). For these images, all competing methods cannot achieve satisfactory results. Some of them even introduce artifacts, such as GDCP [30], Guo et al. [17], and UWCNN_typeII [12]. Additionally, some methods introduce artificial colors. For example, the red object in the fourth image becomes purple and the sand in the first two images becomes reddish. In contrast, our Ucolor not only recovers the relatively realistic color but also enhances details, which is credited to the effective designs of multiple color spaces embedding and the introduction of medium transmission-guided decoder structure. The perceptual scores of our results also suggest the visually pleasing quality of our results.

We present the results of different methods on challenging underwater images sampled from SQUID in Fig. 10. As presented, the input underwater images challenge all underwater image enhancement methods. Ancuti et al. [9] achieves better contrast than the other methods but produces color deviations, e.g., the reddish tone in the third image and the greenish tone in the fourth image. Our Ucolor dehazes the input image, thus improving the contrast of input images. Moreover, our method does not produce obvious artificial colors on the third image. According to the quantitative scores, we can find that the artificial colors significantly affect the perceptual scores given by subjects.

To analyze the robustness and accuracy of color correction, we conduct the comparisons on the underwater Color Checker image in Fig. 11. As shown in Fig. 11(a), the professional underwater camera (Pentax W60) also inevitably introduces various color casts. Both traditional and learning-based methods change the colors of input from an overall perspective. As indicated by the CIEDE2000 values on the results, our Ucolor achieves the best performance (8.21 under CIEDE2000 metric) in terms of the accuracy of color correction.

All the visual comparisons demonstrate that our Ucolor not only renders visually pleasing results but also generalizes well to different underwater scenes.

We first perform quantitative comparisons on Test-S1000 and Test-R90. The average scores of PSNR and MSE of different methods are reported in Table I. As presented in Table I, our Ucolor outperforms all competing methods on Test-S1000 and Test-R90. Compared with the second-best performer, our Ucolor achieves the percentage gain of 20$\%$/59$\%$ and 4.1$\%$/1.3$\%$ in terms of PSNR/MSE on Test-S1000 and Test-R90, respectively. There are two interesting findings from the quantitative comparisons. 1) Although the medium transmission map used in our Ucolor is the same as the traditional GDCP [30], the performance is significantly different. Such a result suggests that the performance of underwater image enhancement can be improved by the effective combination of domain knowledge with deep neural networks. 2) The performance of Unet-U [41] is better than Unet-RMT [41], which suggests that simple concatenation of input image and its reverse medium transmission map cannot improve the underwater image enhancement performance of deep model, and even decreases the performance. 3) The generalization capability of UWCNN models [12] is limited because they require the images to be taken in the accurate type of water as inputs. Due to the limited space, we only present the results of the original UWCNN model that performs the best in Table I in the following experiments, i.e., UWCNN_typeII [12].

Next, we conduct a user study on Test-C60 and SQUID. The average perceptual scores of the results by different methods are reported in Table II, where it can be observed that these two challenging testing datasets fail most underwater image enhancement methods in terms of perceptual quality. Some methods such as Li et al. [4] and UcycleGAN [10] even achieve lower perceptual scores than inputs. For the Test-60 testing dataset, the deep learning-based methods achieve relatively higher perceptual scores. Among them, our Ucolor is superior to the other competing methods. For the SQUID testing dataset, the traditional fusion-based method [9] obtains the highest perceptual score while our Ucolor ranks the second best. Other deep learning-based methods achieve similar perceptual scores. As shown in Fig. 10, all deep learning-based methods cannot handle the color deviations of images in SQUID well, while the haze can be removed. In contrast, the fusion-based method [9] achieves better contrast, thus obtaining a higher perceptual score. Observing the scores of non-reference image quality assessment metrics, we can see that Li et al. [4] obtains the best performance in terms of UIQM and UCIQE scores. For the NIQE scores, our Ucolor achieves the lowest NIQE score on the SQUID testing set while Ancuti et al. [9] obtain the best performance on the Test-C60 testing set.

To demonstrate the robustness to different cameras and the accuracy of color restoration, we report the average CIEDE2000 scores on Color-Checker7 in Table III. For the cameras of Pentax W60, Pentax W80, Cannon D10, Fuji Z33, and Olympus T6000, our Ucolor obtains the lowest color dissimilarity. Moreover, our Ucolor achieves the best average score across seven cameras. Such results demonstrate the superiority of our method for underwater color correction. It is interesting that some methods achieve worse performance in terms of the average CIEDE2000 score than the original input, which suggests that some competing methods cannot recover the real color and even break the inherent color.

We perform extensive ablation studies to analyze the core components of our method, including the multi-color space encoder (MCSE)), the medium transmission guidance module (MTGM), and the channel-attention module (CAM). Additionally, we analyze the combination of the $\ell_{1}$ loss and the perceptual loss. More specifically,

• •

  w/o HSV, w/o Lab, and w/o HSV+Lab stand for Ucolor without the HSV, Lab, and both HSV and Lab color spaces encoder paths, respectively.

• •

  w/ 3-RGB means that all inputs of three encoder paths are RGB images.

• •

  w/o MTGM refers to the Ucolor without the medium transmission guidance module.

• •

  w/ RDCP and w/ RUDCP are the models by replacing the medium transmission map estimated via [30] with the algorithms in [58] and [3], respectively.

• •

  w/o CAM stands for the Ucolor without the channel-attention module.

• •

  w/o perc loss means that the Ucolor is trained only with the constraint of $\ell_{1}$ loss.

The quantitative PSNR (dB) and MSE ($\times 10^{3}$) values on Test-S1000 and Test-R90 are presented in Table IV. The visual comparisons of the effects of reverse medium transmission maps, the contributions of each color space encoder path, the effectiveness of channle-attention module, and the effect of perceptual loss are shown in Figs. 12, 13, 14, and 15, respectively. The conclusions drawn from the ablation studies are listed as follows.

1) As presented in Table IV, our full model achieves the best quantitative performance across two testing datasets when compared with the ablated models, which implies the effectiveness of the combinations of MCSE, MTGM, and CAM modules.

2) Our RMT map can relatively accurately assign the high-quality degradation regions a higher weight than the RDCP map and the RUDCP map, thus achieving better visual quality, especially for high scattering regions as shown in Fig. 12. Additionally, the quantitative performance of the Ucolor is better than the ablated models w/ RDCP, w/ RUDCP, and w/o MTGM, which suggests the importance of an accurate medium transmission map.

3) The ablated models w/o HSV+Lab and w/ 3-RGB produce comparable performance as shown in Table IV and Fig. 13. The results indicate that aimlessly adding more parameters or the same color encoder will not bring extra representational power to better enhance underwater images. In contrast, a well-design multi-color space embedding helps to learn more powerful representations to improve enhancement performance. In addition, removing any one of the three encoder paths (w/o HSV and w/o Lab) will decrease the performance as shown in Table IV.

4) The ablated model w/o CAM produces an under-saturated result as shown in Fig. 14. This may be induced by removing the CAM module that integrates and highlights the features extracted from multiple color spaces.

5) The quantitative results in Table IV show that only using the $\ell_{1}$ loss can slightly improve the quantitative performance on Test-S1000 in terms of PNSR and MSE values. However, from the visual results in Fig. 15, it can be observed that the enhanced image by Ucolor trained with the full loss function (i.e., the combination between the $\ell_{1}$ loss and the perceptual loss) is better than that by Ucolor trained without the perceptual loss. Thus, it is necessary to add the perceptual loss for improving the visual quality of final results. Note that only using the perceptual loss for training does not make sense, and thus we did not conduct such an ablation study.

When facing underwater images with limited lighting, our Ucolor, as well as other state-of-the-arts might not work well. Fig. 16 presents an example where both our Ucolor and latest deep learning-based Water-Net [18] fail to produce a visually compelling result when processing an underwater image with limited lighting. The potential reason lies in few such images in the training data sets. Thus, it is difficult for supervised learning-based networks such as Water-Net and Ucolor to handle such underwater images. The stronger ability and more diverse training data that handle such kinds of underwater images will be our future goal.