USLN: A statistically guided lightweight network for underwater image enhancement via dual-statistic white balance and multi-color space stretch

In the underwater environment, images are rarely color balanced. In general, there will be a phenomenon of higher blue and green light intensity and lower red light intensity. Hence, the first step in our model is to make image balanced. Inspired by white balance algorithm, we propose dual-statistic white balance module.

The structure of dual-statistic white balance module is shown in Fig. 3. This module consists of two trainable balance modules, one is based on Gray World theory, while the other is based on White Patch algorithm. Next, we will explain the detailed derivation process of this module.

Von Kries hypothesis can be used to express the relationship between underwater images and ground-truth images:

Where $I$ denotes pixels values of underwater images, while $I^{\prime}$ denotes those in ground-truth images, and $e_{R}$, $e_{G}$, $e_{B}$ denote scale factors that are applied independently to the three color channels.

However, Von Kries hypothesis considers the R, G, B color channels independently. More generally, we need to consider the relationship between color channels and correction of each channel, so as Fig. 4 shown that we can use pointwise convolution to generalize Von Kries hypothesis as:

Where $Conv_{i}(\cdot)$ means the convolution operation, subscript i means the size of convolutional kernel.

However, it is difficult to learn the relationship between underwater images and enhanced images only by pointwise convolution (Eq. (2)). Inspired by two white balance algorithms, Gray World and White Patch, we add average and maximum of images into Eq. (2) to help model learn this relationship more easily.

The Gray World theory assumes that the average value of color object in a perfect image is gray, which means the average pixels of the three channels R, G, and B tend to the same value. Hence, scale factors of each channel, $e_{R}$, $e_{G}$, $e_{B}$ can be set based on GW theory:

Where $A_{c}$ denotes the average value of $c$ channel in original image.

Therefore, inspired by Eq. (4), we add average of each color channel in Eq. (2), which can be expressed:

where $\overline{A}=[\frac{1}{A_{R}},\frac{1}{A_{G}},\frac{1}{A_{B}}]\in\mathbb{R}^{3\times 1}$, $\otimes$ denotes pixel-wise multiplication, $I^{GW}$ means image enhanced by the trainable white balance module based on Gray World algorithm. Gray World algorithm can be considered a special case of Eq. (5).

White Patch algorithm assumes the maximum response of the RGB color channel is caused by the white patch in the scene. Theoretically, the white patch can reflect the color of the scene light, so the largest value in the RGB channel will be used as the light source. According to this hypothesis, scale factors of each channel can be expressed as:

Where $M_{c}$ denotes the maximum value of $c$ channel in original image.

Same as before, inspired by Eq. (7), we add maximum of each color channel in Eq. (2), which can be expressed:

where $\overline{M}=[\frac{1}{M_{R}},\frac{1}{M_{G}},\frac{1}{M_{B}}]\in\mathbb{R}^{3\times 1}$, $I^{WP}$ means image enhanced by the trainable white balance module based on White Patch algorithm. White Patch algorithm can be considered a special case of Eq. (8).

Finally, two enhanced image are merged together by going though 3×3 convolutional layer and pixel-wise adding together.

where $I_{final}$ means the final result enhanced by dual-statistic white balance module. $\oplus$ means pixel-wise addition.

Due to the low histogram range, underwater images often suffer from low contrast and visibility. Therefore, histogram stretching is adopted to provide a wider pixel distribution of the image channels to improve the image contrast. The following equation is used for linear contrast stretching:

Where the $P_{in}$ and $P_{out}$ are the input and output pixels intensity values, respectively, $i_{max}$ and $i_{min}$ represent the maximum and minimum intensity value of the input image, respectively, $o_{max}$ and $o_{min}$ represent the maximum and minimum intensity value of the desired output image, respectively.

In histogram stretching-based methods, $i_{max}$ and $i_{min}$ are easy to get, but $o_{max}$ and $o_{min}$ must be defined artificially. For example, Kashif Iqbal [3] used global histogram stretching to enhance the contrast of the image, setting $o_{max}$, $o_{min}$ to 1, 0. However, such a simple setup may not achieve good results. We can use pointwise convolunsion to make Eq. (10) trainable by next steps:

Where $\widetilde{I}=[\widetilde{I}_{R},\widetilde{I}_{G},\widetilde{I}_{B}]\in\mathbb{R}^{3\times H\times W}$, $I^{r}$ denotes the histogram stretched pixel value in RGB color space. Global histogram stretching can be considered a special case of Eq. (12).

Fig. 5 presents the details of the trainable stretch module in RGB space.

In addition to operating in the RGB color space, inspired by traditional algorithms [13, 3, 26] and Ucolor [8], we also stretch images in HSI and Lab spaces. HSI color space reflects the way the human visual system perceives color, which intuitively reflects the hue, saturation and intensity. And Lab color space is designed to approach human perceptual vision, which contains almost every shade of every color in nature. As Fig. 6 shown that the same image has different color distribution in various color spaces. Thus, the multiple color spaces stretch can facilitate the adjustment of underwater images color. In addition, the transformation relationship of images in these color spaces is not linear [27], which enables the deep-learning model to learn more complex mapping relationships. We also design ablation study to verity the contribution of multi-color space stretch operation.

Concretely, we convert the image into HSI and Lab color spaces and then stretch the S, I channels and L, a, b channels, respectively. Note that we do not stretch H channel, because our ablation studies verity that changing the value of the H channel will damage the color of the image. Interestingly, this conclusion is consistent with the statistic-based methods like [3], [13] and [28].

At last, three stretched images are converted back to RGB color space and merged together by going though 3×3 convolutional layer and pixel-wise add up.

Where $I^{r}$, $I^{h}$, $I^{l}$ denote the histogram stretched pixel value in RGB, HSI, Lab color spaces, respectively.

It should be noted that we made a pytorch implementation that converts image between RGB, HSV, and Lab color spaces allowing differentiable back-propagate, so our model is trained in an end-to-end manner.

Drawing on deep learning techniques, we introduced residual enhancement module to our model, which aims to preserve the data fidelity and increase capacity of model [29]. Fig. 7. shows the structure of residual-enhancement module, which consists of a convolutional layer and an activation function.

where $z$ and $res$ mean the input and output of residual-enhancement module, respectively, $\delta(\cdot)$ represents the Tanh activation function.

The kernel size of convolutional layer is set as 3×3, because the main convolutional layers in dual-statistic white balance module and multi-color space stretch module are pointwise convolution, which ignores the relationship between adjacent pixels. Therefore, the 3×3 convolutional layers in residual-enhancement module can alleviate this problem. For activation function, we chose Tanh activation function, because the output of Tanh activation is between -1 and 1, which means that residual-enhancement module can choose to reduce or increase the intensity of the enhanced pixels based on the features.

Inspired by Hang Zhao’s work [30] and Johnson’s work [31], we used combined loss to train our model, which consists of MAE loss, SSIM loss and VGG perceptual loss.

MAE loss is to calculate the mean of the absolute difference between all pixels of the output image, $\hat{y}$, and the label, $y$, it can be expressed as:

SSIM Loss is to improve the structural similarity between the feature representations of the output image, $\hat{y}$, and the label, $y$:

where u and $\sigma$ denote the mean, standard deviation, and covariance of images. $\mu_{1}$ and $\mu_{2}$ are used to stabilize the division.

VGG perceptual loss is based on pretrained VGG networks. Let $\phi(\cdot)$ be the last convolutional layer of VGG-19 networks [27]. The perceptual loss is calculated the MAE loss of the feature representations of output image, $\hat{y}$, and the label, $y$, assuming that the total number of pixels is $M$:

The final combined loss is the linear combination of the MAE loss, SSIM loss and VGG perceptual loss:

where $\lambda_{1}$ and $\lambda_{2}$ are empirically set to 0.25 and 1, respectively, for balancing the scales of different losses.

The validity of our combined loss function will be demonstrated in later ablation experiments.

From perspective of attention mechanism, we think that our model can be seen as a special case of channelwise attention networks. Classic channelwise attention networks, SE block [32] can be roughly divided into three stages: squeeze, excitation and scale. SE block first squeezes global spatial information into a channel descriptor, then the global information is embedded by channelwise multiplying the excited channel descriptor and feature map together. In our model, the squeezed channel descriptors include average, maximum, and minimum of images, then they are excited by Eq. (5), Eq. (8), and Eq. (12), respectively. These operations can exploit the global information effectively.

In our work, we adopted four benchmarks for evaluation, Uimagine [9], UIEB [6], UFO [33], and Color-Check7 [34]. The scale of these datasets is various, so that we can compare these methods comprehensively.

In UIEB [6], Uimagine [9], and UFO [33], there are many paired images and unpaired images which were divided as shown in Tab. II.

Color-Check7 [34] contains seven original underwater images with Color Checker. These images are taken by seven different cameras, Canon D10, Fuji Z33, Olympus Tough 6000, Olympus Tough 8000, Pentax W60, Pentax W80, and Panasonic TS1, denoted as D10, Z33, T6000, T8000, W60, W80, TS1. Considering there is no corresponding train set in Color-Check7, we got three results from models trained on Train-800, Train-3145 and Train-1500, respectively. Finally, the best score is used to evaluate the performance.

For test sets in UIEB, Uimagine and UFO, we chose three full-reference evaluation metrics and one no-reference evaluation metric. Concretely, 1) The full-reference evaluation metrics consist of Peak Signal-to-Noise Ratio (PSNR) [35], Structural Similarity (SSIM) [35] and Mean Squared Error (MSE), which are used for paired test sets (Test-90, Test-555, Test-120). A higher PSNR or a lower MSE score means that the output image and the label image are closer in content, while a higher SSIM score means that the two images are more structurally similar. 2) The no-reference evaluation metric is UIQM [36], which is used for unpaired test sets (Test-60, Test-1270). A higher UIQM score suggests a better human visual perception. UIQM is the linear combination of UICM, UISM, and UIConM, which reflect the degree of color deviation, brightness, and contrast, respectively. However, it is worth mentioning that these no-reference metrics cannot accurately reflect the quality of an image in some cases [6] [37], so scores of UIQM (UICM, UISM, UIConM) are only provided as references for our study. We will present enhanced unpaired images in Visual Comparisons section for readers to score.

For Color-Check7, each colorcheck in Color-Check7 consists of 24 color patches. We used CIEDE2000 [38] to evaluate the enhanced results. Concretely, we cropped the 24 color patches for each enhanced result and measured the average value of each color patch, then calculated the perceptual differences between corresponding color patches of ground-truth Color Checker and those of enhanced results.

We used python and pytorch framework via NVIDIA RTX3080 on windows 10. Adam was used for network optimization, while the initial learning rate was set as 0.01 with 5% reduction per epoch. We trained our model for 100 epochs and the batch size was set as 10.

Different from fixed size (256×256) of images in Uimagine, the size of images in UIEB is unfixed. In order to use batch learning, all images in UIEB were resized to 256×256. For UFO, the raw images are low-resolution (320×240), while the ground truth images are high-resolution (640×480). However, our research in this paper has nothing to do with the super-resolution task, so all images in UFO were also resized to 256×256.

Note that except for resizing images, we did not pre-process the images in any other way in our experiments.

We selected eight works to make fair comparisons with our model, including two statistic-based methods (UCM [3], RGHS [4]), two physical physical-model-based methods (Li et al. [12], ULAP [5]), four deep-learning-based methods (Waternet [6], Ucolor [8], UWCNN [7], FUnIE-GAN [9]). UCM integrates White Balance and histogram stretch to improve image quality, which is highly similar to our model. The comparison with UCM aims to demonstrate the advantages of our trainable modules.

For these traditional unsupervised methods, we used them directly on the test sets. For these deep-learning-based methods, we used the code and training strategy provided by the authors. Note that although FUnIE-GAN can be trained with unpaired data, we also trained this model in paired train sets to ensure the fairness of the experiment.

In addition to the quality of the enhancement results, for deep-learning-based models, GFLOPs and the amount of parameters are also important evaluation metrics. As shown in Table III, our model takes the least parameters and GFLOPs. Compared with the second-lightest model, the amount of parameters and GLOPs of our model are only 0.22% of UWCNN (amount of parameters) and 7.89% of FUnIE-GAN (GFLOPs).

We first make a quantitative comparison on Test-90, Test-555, and Test-120. The average scores of MSE, PSNR, and SSIM of different methods are shown in Table IV. As shown in Table IV, although there are only few parameters in our model, our method still performs better than the complex deep-learning-based models. Our model achieves the lowest MSE score and highest PSNR, SSIM scores in all three datasets. In addition, we also found that 1) although our method principle is roughly the same as UCM model, which is based on white balance and histogram stretch, our model performs much better than UCM, which suggests the superiority of our trainable modules. 2) Our model shows excellent performance on small datasets UIEB. Compared with the second-best performer, our USLN achieves the percentage gain of 16%, 8% and 5% in terms of MSE, PSNR and SSIM, respectively. This means that our model can learn the mapping of underwater images and enhanced images well with small dataset.

Then, for Color-Checker7, we show the average CIEDE2000 scores of different methods in Table V. Our model achieves highest color similarity in Z33, T6000 and T8000. Moreover, our model achieve the lowest average CIEDE2000 score (10.71 under CIEDE2000 metric). These results show that our model is robust to different cameras and can restore the real color of underwater images with high accuracy. Interestingly, many compared methods get worse average CIEDE2000 score than raw images, which suggests that these methods even destroy the images color.

Next, we use UIQM as evaluation metric for unpaired test set (Test-60, Test-1270). The average scores of different methods are shown in Table V. As Table VI shown, Our model gets highest UICON in Test-90, and second best UICM and UICON scores in Test-1270. There are two interesting findings from the Table V, 1) Different from full-reference evaluation metrics, deep-learning-based models have not achieve much better results than traditional methods in no-reference evaluation metrics. Conversely, for UICM, Li et. al and UCM rank the best and second best respectively in both datasets. 2) Although FUnIE-GAN and Ucolor perform best in Test-90 and Test-1270, respectively, they do not perform well in another dataset. In comparison, our model achieves robust performance in both two test sets.

In this section, we conduct several visual comparisons on testing datasets.

We first present raw images and enhanced images in paired datasets in Fig. 9. Compared with other SOTA models, our result is closest to the ground-truth images and achieves the highest PSNR and SSIM scores. Note that, as table.1 shown, traditional methods get much lower PSNR/SSIM scores than deep-learning-model, so due to length limitation, we only present images enhanced by deep-learning-based models in Fig.7.

The underwater images suffer from different color deviations are shown in Fig. 9LABEL:sub@fig_first_case. The first two images (left) are only distorted slightly, while the other images are more distorted. For slightly distorted images, all compared methods fail to restore the complete scene structure, these methods all get even lower PSNR or SSIM than raw images. In comparison, our model improves the quality of these images effectively without over-enhancement. For the rest images which are deviated by greenish and blueish color. The compared models either fail to remove the greenish and blueish color completely or introduce undesirable color artifacts. For example, these images enhanced by FUnIE-GAN and Ucolor are greenish and reddish, respectively. In comparison, our model recovers the color of image and achieves the best scores.

To demonstrate the robustness of our model, we then present the enhanced Color Checker images of different methods. The raw image shown in Fig. 8 is taken by Panasonic TS1, which suffers from severe red color cast. All compared methods cannot achieve satisfactory results, some of them even reduce the quality of image. According to CIEDE2000 metric, our method gets the best score (9.90) and restores the most accurate colors of the raw image.

We then show raw images and enhanced images in unpaired datasets, Test-60 and Test-1270 in Fig. 10, Fig. 11, and Fig. 12. Noted that we only use UIQM as the reference to evaluate these images, the accurate perceptual quality of these enhanced images need the readers to score.

For Fig. 10, the raw image sampled from Test-1270 suffers from serious blueish deviation. Our result gets the best UIQM. Compared with SOTA models, our model successes to recover the natural color and structural details.

For Fig. 11 and Fig. 12 which are sampled from Test-60, these two images are much more challenging. As Fig. 11 shown, the obvious light limitation of the raw image results in low contrast. The compared models tend to make the image darker, some models even introduce reddish color. In comparison, our model increases both brightness and contrast, making the details of the image clear. Fig. 12LABEL:sub@input suffers from obvious yellow deviation, all compared methods cannot achieve satisfactory results. For example, although Fig. 12LABEL:sub@waternet which is enhanced by Waternet achieves the best UIQM, it is obvious that the left part of the image is over-enhanced, introducing unwanted blue, while the right part is under-enhanced, failing to remove yellow completely. In comparison, our model removes the yellowish color and makes the color balanced.

We perform extensive ablation studies to analyze the advantage of the core modules, including residual-enhancement module (REM), multi-color space stretch module (MCSM), and dual-statistic white balance module (DSBM). In addition, we analyze the combined loss in our model. More specifically,

• •

  w/o Lab+SI, w/o Lab, and w/o SI mean that the multi-color space stretch module without both Lab and SI color channels, Lab channels, and SI channels, respectively. w/ H means that we do not separate H channel out in multi-color space stretch module.

• •

  w/o avg and w/o max, and w/o max+avg stand for dual-statistic white balance module without average, maximum, and both average and maximum, respectively.

• •

  w/o RES refers to USLN without residual-enhancement module, while w/o Tanh refers to the Residual-Enhancement Module without Tanh activation function.

• •

  w/o VGG+SSIM, w/o SSIM, and w/o VGG mean that USLN is trained with MAE loss, MAE+SSIM loss, and MAE+VGG loss, respectively.

We used MSE, PSNR, and SSIM to evaluate the results on Test-90, which are shown in Table VI. Compared to the ablation models, full model achieves the best quantitative performance on all metrics. The visual comparisons of the contributions of dual-statistic white balance module, muti-color stretch model, and residual-enhancement model are presented in Fig. 13, Fig. 14, and Fig. 15, respectively. The conclusions drawn from the ablation study for the each module are listed as follows:

1)The ablated models w/o avg and w/o max produce comparable performance as shown in Table VII and Fig. 13. Dual-statistic balance module help model to balance the images base on two representative values, average and maximum, while removing either of them will degrade performance as Table VII shown. As Fig. 13 shown, the full model removes the color casts much more completely than ablated models.

2)As Table VII shown, the lack of stretch in either HSI or Lab space will decrease performance. In addition, as Fig. 14 shown, the full model enhances images from three color spaces, which is effective to restore the natural color of image. Interestingly, we find blindly add all channels in HSI color space will reduce the performance, while adding only S and I channels can improve performance effectively.

3)Residual-enhancement module can preserve the image fidelity, making the enhanced images look more natural. As presented in Fig. 15, w/o Res model over-enhances the image, making aquatic weeds be greyish, while full model restores normal green color of aquatic weeds. Moreover, the lack of Tanh activation function will degrade the performance as Table VII presented.

4)Ablation study towards the combined loss is presented in Table VII. SSIM loss in the combined loss function can improve the structural similarity of the output image as shown in Table VII, while VGG perceptual loss can improve the visual quality of output. By adding perceptual loss and SSIM loss to the MAE loss, the visual quality of final result is improved.

Enhancing the underwater image with severe color cast is very challenging. Both our model and other SOTA models fail to work well. As Fig. 16 shown, although our model and Ucolor remove the yellow cast, it is far from achieving the desired visual quality. Ucolor makes image much darker, while USLN introduces unwanted red. The main reason is the lack of similar images in train sets, limiting the ability of deep-learning-based methods to deal with such images.