WaterFlow: Heuristic Normalizing Flow for Underwater Image Enhancement and Beyond

In recent years, numerous underwater image enhancement methods have been proposed. The early traditional methods often adjust the pixel distribution of different color channels to weaken the degradation of natural light underwater. Ghani et al. (Ghani and Isa, 2015) extended the histogram of the red channel and the blue channel of the underwater image upward and downward respectively according to the law of Rayleigh distribution. Li et al.  (Li et al., 2016) corrected the red channel according to the Gray-World assumption theory after the color restoration of the blue-green image through the dark channel. Although these methods can weaken the fading of light on different color channels, it is easier to introduce underwater images into artificial colors due to the lack of guidance from physical models. Therefore, model-based methods (Li et al., 2017; Wang et al., 2017; Peng et al., 2018; Yang et al., 2023) are also widely used to improve the interpretability of networks by incorporating domain-specific prior knowledge.

In recent years, many methods based on deep learning have been proposed: Li et al. (Li et al., 2021a) introduced the multi-color space into the transmission-guided network to reduce the influence of color casts on underwater images. Mu et al. (Mu et al., 2022) introduced a bi-level model that hierarchically incorporates various knowledge modalities to enhance the quality of underwater images. Jiang et al. (Jiang et al., 2022) proposed a perceptual adversarial mechanism and introduced the global module to narrow the gap between enhanced images and reference images. Huang et al. (Huang et al., 2023) proposed a semi-supervised framework based on the mean-teacher approach to enhance the generalization ability of the enhance model on real-world data.

Normalizing flow facilitates the transformation between complicated probability distribution and Gaussian distribution through bijection functions and differentiable mappings. The normalizing flow has received far less attention than GANs and VAEs because it requires a significant quantity of video memory with a fine structural design to handle a large amount of reversible computing. Many methods for accelerating the calculation have been proposed in recent years. Dinh et al. (Dinh et al., 2014) proposed the additive coupling layer, which makes the flow easier to calculate the determinant of Jacobian matrix. In order to increase the log-likelihood, Kingma et al. (Kingma and Dhariwal, 2018) developed the actnorm to normalize each channel of input features and $1\times 1$ convolution in place of the permutation (Dinh et al., 2014). Conditional normalizing flow was introduced for increasing the expressiveness and flexibility of normalizing flows, and has been widely used in vision tasks (Lugmayr et al., 2020; An et al., 2021; Abdal et al., 2021; Li et al., 2021b; Wang et al., 2022; Jiang et al., 2023b). However, there is little precedent dedicated to adopting flow-based methods to solve the problem of ill-posed underwater image enhancement tasks.

In recent years, due to the rapid development of deep learning, the effect of object detection has also been significantly improved. The existing methods are generally divided into one-stage method and two-stage method. The two-stage method (Girshick et al., 2014; Ren et al., 2015; He et al., 2017; Sun et al., 2021; Cheng et al., 2022) first needs to extract the region of interest proposals area, and then classify it through the classification of neural network. By contrast, the classification label and boundary information are often directly obtained by the one-stage method (Redmon and Farhadi, 2018; Kim and Lee, 2020; Feng et al., 2021; Xie et al., 2022).

The current detectors can achieve good performance when facing in-air images. However, the problems of scene blur and light imbalance faced by underwater images will greatly affect the effectiveness of the detectors. Existing studies (Liu et al., 2021, 2022a; Ma et al., 2022; Liu et al., 2023b, a) have used low-level visual enhancement modules as preprocessing steps for detection tasks. However, the existing enhanced networks for the underwater image often do not play a significant role in the detector, which ignores the position and semantic information while restoring the contrast of the underwater image.

Hybrid Invertible Block (HIB) was introduced as the main part of the heuristic normalizing flow. It ensembles the heuristic prior into the data-driven network to build the invertible mapping relationship between the underwater image and its clear counterpart. As illustrated in Fig. 2, the underwater image restoration and the underwater image degradation process are combined in a invertible mapping manner during the training process. In the forward process, the underwater image is squeezed and input into multiple invertible blocks to generate enhanced images. In the reverse process, clear images are inversely input to invertible blocks with shared parameters to generate degraded images.

The Hybrid Invertible Block consists of the following parts: Actnorm (Kingma and Dhariwal, 2018), Conv 1 $\times$1 (Kingma and Dhariwal, 2018), Heuristic Prior Injector, Conditional Affine Coupling (Li et al., 2021b) and squeeze operation.

Actnorm (Kingma and Dhariwal, 2018) is a normalizing operation, which can change the input tensor into the zero mean and unit variance tensors. Conv $1\times 1$ (Kingma and Dhariwal, 2018) can convolve the input tensor of $c\times h\times w$ by the weight matrix of $c\times c$ to make the calculation of Jacobian matrix and network inversion easier.

Heuristic Prior Injector was designed to better characterize the underwater imaging model by embedding a heuristic prior, which will be introduced in detail in section 3.2.

Conditional Affine Coupling was proposed by (Li et al., 2021b), which realizes the mutual conversion with conditional information between output and input through reversible transformation. We modified the previous formulation for adjusting the underwater image enhancement framework, which is described as:

where $\mathbf{u}_{i+1}=\mathcal{C}(\mathbf{u}_{i+1}^{1},\mathbf{u}_{i+1}^{2})$, $\mathcal{C}$ represents the concatenation operation. $\phi_{i}$, $\rho_{i}$, $f^{A}_{i}$,  and $f^{B}_{i}$ denote convolutional networks. $i\in\{1,...,n-1\}$ represents the $i$-th flow steps, $I_{u}$, $I_{g}$ and $I_{d}$ will be explained in detail in section 3.2. Squeeze operation can convert the image of $c\times h\times w$ to $4c\times\frac{h}{2}\times\frac{w}{2}$. Unsqueeze is the reverse operation of squeeze. By the combination of these operations, the proposed architecture can attain accurate mapping between the underwater image and it clear counterpart while maintaining reversibility.

The Heuristic Prior guided Injector (HPI) is designed to incorporate the estimated physical imaging parameters into the invertible block. According to the underwater image formation model (Drews et al., 2013; Chiang and Chen, 2011), the enhanced image can be represented as:

where $I$ and $J$ denote the underwater image captured by sensors and the enhanced image at pixel point $x$ respectively. $c\in\{r,g,b\}$ represents the corresponding color channels. $B$ and $t$ indicates the ambient light and the medium transmission coefficient. $t$ can also be described as $t(x)=e^{\beta d(x)}$ with the scene depth $d(x)$ and the attenuation coefficient $\beta$ according to the Beer-Lambert law (Bouguer, 1729).

Considering that the gradient map provides information about the edges and contrast in the image, which can be used to estimate the scatteration of light by the atmosphere and fine impurities in the water. Furthermore, the depth map provides information about the attenuation of the light propagation, which is also a key factor in the Beer-Lambert law. Inspired by (Peng et al., 2018), we first estimate the depth map $I_{d}$ and the gradient map $I_{g}$ as auxiliary information from the underwater image.

Then, we concatenate them with underwater image as input to the Heuristic Prior guided Encoder (HPE) $\mathcal{H}$ to progressively estimate the imaging parameters by extracting the depth, gradient and color information of the underwater image. The equation can be illustrated as:

where $I_{u}$ represents underwater images. $\mathcal{C}$ and $\mathcal{S}$ denote the concatenate and split operation respectively. $T$ is the reciprocal of $t$. The detailed architecture of $\mathcal{H}$ is in the appendix due to limited space. After estimating $T$ and $B$, we insert them into the HIB as heuristic information. The operation can be formulated as:

where $\mathbf{u}_{i}$ represents the middle feature of the $i$-th flow steps, $\odot$ denotes dot product. By introducing the underwater image formation model based heuristic prior to the invertible blocks, the proposed method can generate enhanced images, which are more suitable for realistic underwater scenes.

Object detection aims to acquire the location and category of each object. Hence, the performance of object detection will be improved if the enhanced image holds a more profound level of semantic feature and object localization information. Inspired by the perceptual loss (Johnson et al., 2016), deep neural networks designed for high-level visual tasks can retain semantic information and own the ability to describe potential features. Therefore, in order to make the enhanced images potentially improve the effectiveness of the underwater object detection, we propose a Detection Perception Module (DPM) to improve the performance of the enhancement module in preserving and extracting detection-oriented perceptual features.

In the whole training process, we first use the enhancement module as the data preprocessing stage of underwater object detection. After that, the enhanced image was directly fed into the DPM and conducted separate detection training. Then, we jointly optimize the pre-trained DPM and enhancement module through different benchmarks. Specifically, the underwater image is first input to the enhancement module to obtain enhanced data. We input it into the detection module to transmit the high-level implicit perceptual features extracted from the network to the enhancement module, so as to guide the visual improvement effect to be more conducive to subsequent object detection tasks.

Contrastive learning has been applied to multiple low-level visual tasks (Liu et al., 2022b, c, 2023c; Jiang et al., 2023a). We introduce contrastive learning to improve the quality of enhanced images by making them closer to the in-air images and more isolated from the underwater images. We have designated the reference image as a positive sample and the original underwater image as a negative sample. VGG19 (Simonyan and Zisserman, 2014) was used to extract the implicit characteristic of enhanced images. The contrastive loss is expressed as follows:

where $\mathcal{V}_{i}$ represents the $i$-th layer of VGG19 (Simonyan and Zisserman, 2014), and $\rho_{i}$ denotes the $i$-th weight of each layer. $G_{E}$ represents the proposed enhancement network. $I_{u}$ and $I_{r}$ represent the underwater image and the reference image respectively.

We additionally incorporate the style loss (Deng et al., 2022) to make the generated output closer to the style pattern of the reference image. The definition of style loss function $\mathcal{L}_{s}$ is described as follows:

where $\mu$ and $\nu$ represent the mean and variance of the image. $N$ is the number of layers.

Localization loss $\mathcal{L}_{\text{loc}}$ and classification loss $\mathcal{L}_{\text{cla}}$ are used as detection-driven loss $\mathcal{L}_{\text{det}}$, which is illustrated as:

where classification loss is designed to minimize the discrepancy between the prediction category and the ground truth category. The localization loss is adopted to reduce the position difference between the prediction box and the ground truth box. We employ Focal loss (Lin et al., 2017) and GIoU loss (Rezatofighi et al., 2019) as classification loss and localization loss respectively.

Ideally, the incorporation of the reverse mapping process could impose a regular constraint on the degraded image, thereby enhancing the performance of the forward mapping (Guo et al., 2020). Therefore, L1 loss is used as a bilateral constraint to make the output of the forward and the reverse process closer to the reference and underwater image respectively. The formula can be described as:

where $G_{E}^{-1}$ denotes the reverse transformation of $G_{E}$. Therefore, the total loss of network training is expressed as follows:

Our network is implemented by the pytorch and trained on NVIDIA RTX 3090 GPU. We first randomly crop the UIEBD (Li et al., 2019a) dataset into the size of $384\times 384$ and use the cropped dataset as the training dataset for underwater image enhancement. UCCS (Liu et al., 2020) and Aquarium (Prytula, 2020) datasets are used as the training dataset for underwater object detection. We first train the heuristic normalizing flow and Detection Perception Module separately for $5\times e^{5}$ iterations. Then, we jointly train the pre-trained modules for $3\times e^{5}$ iterations. During the training process, Adam was used as our optimizer with the uniform learning rate of 1e-6 and the batch size of 2. $\lambda_{1}$, $\lambda_{2}$, $\lambda_{3}$, $\lambda_{4}$ are set to $1$, $100$, $0.1$, $1$.

For visual comparison, Fig. 3 shows the enhancement results of all methods on UIEBD dataset. Water-Net and Ucolor not only fail to recover the original reflection, but also suffer from low contrast and saturation. DLIFM, TOPAL, TACL and SemiUIR have limited effects on underwater image enhancement, which still remain conspicuous color deviation at close scenes. Compared with the other methods, the proposed method recovers the degradation of light and rectifies the unexpected coloration. In addition, we demonstrate the pixel distribution similarity diagram of a particular region between the corresponding enhanced and reference images. The red line and the blue line in the diagram indicate the pixel distribution of the reference image and the enhanced image respectively. Notably, the proposed method yields enhanced results closest to the reference image.

The comparison on the EUVP dataset is shown in Fig. 4. Ucolor and TOPAL introduce artifacts and unnatural colors. Water-Net, TACL and SemiUIR fail to alleviate the scattering of light underwater. Although DLIFM produces pleasing images, the instructive details of the image have not been recovered well. Compared with the other methods, the proposed method not only effectively corrects the color deviation, but also restores the scene radiance and contrast of the image. Fig. 4 also shows the distribution of RGB color space for the images obtained by different methods. In general, the proposed method more effectively addresses the swift attenuation of red wavelengths underwater, with its color space distribution closely approximating that of in-air images.

The qualitative comparison on the U45 datasets is shown in Fig. 5. It is obvious that the red light and blue light decay significantly more than the green light underwater in this acquisition environment. However, the other compared methods exhibit limited effectiveness in restoring the decay of blue light. The boxplot for the intensity distribution of the RGB color channels is set below each image, with the horizontal axis representing pixel intensity. Additionally, semi-transparent rectangles are introduced to expound the relationship between distinct color spaces more lucidly, with the gray rectangle signifying the region where outliers emerged. In comparison with the proposed method, the other methods not only have a noticeably inferior ability to revive the attenuation of blue light but also introduce certain aberrant outliers, adversely impacting the efficacy of the enhanced images.

The visual comparison on the UCCS datasets is shown in Fig. 6. TOPAL fails to restore the color deviation of the image, which suffers from low saturation. Unreasonable color artifacts are introduced by Water-Net, TACL and SemiUIR. Informative details are ignored by DLIFM and Ucolor, resulting in obscure reflections and low contrast. Compared with other methods, our method significantly restores the brightness and contrast of underwater scenes without the interference introduction.

We utilize the same comparison methods to assess whether the proposed enhancement results implicitly contain more semantic features that are applicable to subsequent detection tasks. The visual comparison of UCCS dataset is shown in Fig. 7. The proposed method achieves a higher detection confidence than other methods with more detected trepangs and urchins. The qualitative comparison of the Aquarium dataset is shown in Fig. 8, which clearly shows the remarkable advantages of the proposed method in adapting to subsequent detection tasks.

The quantitative results of representative methods for underwater image enhancement are shown in Tab. 1. It can be observed that the proposed method exceeds all other methods in UICM, and achieves competitive results in UCIQE and UIQM. However, as noted by (Berman et al., 2021; Huang et al., 2023), although both UCIQE and UIQM metrics can represent the degree of restoration of underwater images, they are somewhat heuristic and have limited applicability on evaluating the performance of underwater image enhancement. We further compared fully-reference evaluation on the UIEBD dataset which contains reference images. The proposed method achieves a significant advantage on PSNR and is only lower than SemiUIR on SSIM. However, as noted by (Li et al., 2021a), the reference images of UIEBD are synthetic by multiple enhancement methods, which indicates that some unreliable samples may affect the evaluation of the performance. Combining with NR-IQA and full-reference benchmarks, our method still achieves great advantages in underwater image enhancement, consistent with the qualitative results. The quantitative comparison for underwater object detection is shown in Fig. 9. It can be observed that the proposed method outperforms the other competitive methods in terms of detection accuracy in both UCCS and Aquarium datasets. It can be proven that the proposed results can not only achieve good visual enhancement effects, but also potentially contain more perceptual information which is beneficial for subsequent detection tasks.