UWGAN: U NDERWATER GAN FOR R EAL - WORLD U NDERWATER C OLOR R ESTORATION AND D EHAZING
Anonymous authors
Paper under double-blind review
A BSTRACT
In real-world underwater environment, exploration of seabed resources, underwater archaeology, and underwater fishing rely on a variety of sensors, vision sensor is the most important one due to its high information content, non-intrusive, and passive nature.
However, wavelength-dependent light attenuation and back-scattering result in color distortion and haze effect, which degrade the visibility of images.
To address this problem, firstly, we proposed an unsupervised generative adversarial network (GAN) for generating realistic underwater images (color distortion and haze effect) from in-air image and depth map pairs based on improved underwater imaging model.
Secondly, U-Net, which is trained efficiently using synthetic underwater dataset, is adopted for color restoration and dehazing.
Our model directly reconstructs underwater clear images using end-to-end autoencoder networks, while maintaining scene content structural similarity.
The results obtained by our method were compared with existing methods qualitatively and quantitatively.
Experimental results obtained by the proposed model demonstrate well performance on open real-world underwater datasets, and the processing speed can reach up to 125FPS running on one NVIDIA {+
+} GPU.
I NTRODUCTION
In recent years, underwater vision plays an important role in a lot of different applications.
Therefore, underwater image processing has received extensive attention and research due to the poor underwater imaging environment and image quality.
The main reason is the scattering and attenuation of light, the scattering results in haze effect, and the attenuation of light leads to color cast.
So far many image enhancement algorithms have been proposed, such as white balance algorithm (Liu Y C, 1995), gray world algorithm (Rizzi A, 2002), histogram equalization (Pizer S M, 1987) and fusion algorithm (Ancuti C, 2012), however, these methods are not based on the underwater physical imaging model, so it is challenging and ineffective to apply these algorithms to different underwater scenes directly.
Many underwater image enhancement algorithms based on imaging models have been proposed.
For instance, He et al (He K, 2010) proposed a dark channel prior (Dark channel prior, DCP) dehazing algorithm based on many experiments.
Chiang et al (Chiang J Y, 2011) apply DCP model on underwater image dehazing problem.
These traditional methods are not intelligent, it is very time-consuming to calculate the characteristics of the image.
In these years, the deep learning network developed rapidly, especially the convolutional neural network (CNN), which is used in image classification (Krizhevsky A, 2012), object detection (Redmon J, 2016), and motion recognition (Kuehne H, 2011), the performance is much better than traditional methods.
However, the current research on underwater image enhancement using CNN is limited due to lack of underwater datasets.
It is difficult to obtain images without water in real-world underwater scenes.
Therefore, using synthetic underwater datasets is an important approach (Anwar S, 2018; Ancuti C, 2016; Uplavikar
P, 2019). Some model based on generative adversarial network (GAN (Goodfellow I, 2014)) are used to generating realistic underwater images.
For instance, CycleGAN (Zhu J Y, 2017) generates images through style transfer.
WaterGAN (Li J, 2017) takes in-air images, depth maps and noise vectors as input, followed by a camera model, then output synthetic images.
Based on our experimental results, the image generated by WaterGAN suffers color noise and they differ a lot from real world underwater images.
Therefore, to generate realistic underwater images with both color cast and haze effect, we improved the underwater imaging model, and proposed an unsupervised GAN based on this model to generate realistic underwater images from clear in-air images.
Then, U-Net with different loss functions (Ronneberger O, 2015) is trained to enhance underwater images through synthetic datasets.
Finally, the performance of the proposed algorithm is validated on real underwater images as well as underwater target detection datasets for both low-level and high-level computer vision tasks.
The experimental results show that the proposed method can recover the underwater image while maintaining structural similarities.
Apart from this, the effects of different loss functions in Unet are compared, the most suitable loss function for underwater image restoration is suggested based on the comparison (This part can be found in A PPENDIX ), which provides a new idea for underwater image enhancement.
O UR P ROPOSED M ETHOD
To generate the realistic underwater images (color casts, low contrast and haze effect), we improved underwater imaging model, and proposed an underwater generative adversarial network (UWGAN), which takes in-air RGB-D images and a sample set of underwater images of a specific survey site as input to train a generative network adversarially.
These synthetic underwater images, which were used to train a restoration network based on U-Net (Ronneberger O, 2015) that can enhance underwater images in real-time.
2.1 I MPROVED UNDERWATER IMAGING MODEL
As is well known, a simplified underwater imaging model is shown in Eq.1.
[Equation]
(Eq. 1) where, 𝐼𝐼 ( 𝑥𝑥 ) is the light intensity of each pixel 𝑥𝑥 .
𝐽𝐽 ( 𝑥𝑥 ) is the initial irradiance that not propagating through the water.
𝑇𝑇 ( 𝑥𝑥 ) is the transmission map of the scene. 𝐴𝐴 is the atmospheric ambient light of the scene.
𝛽𝛽 is attenuation coefficient of light of different wavelengths 𝜆𝜆 , and 𝑑𝑑 ( 𝑥𝑥 ) is the range between the scene and the camera.
We can generate underwater-style images using the in-air image and its depth map by Eq. 1, which can well simulate color cast caused by light attenuation in water. However, it is difficult to simulate the haze effect caused by the scattering of water impurities.
As shown in Figure 4, obvious haze effect can be observed on real underwater images.
Inspired by related dehazing methods (Ancuti C, 2016), we improved the second term in Eq. The improved imaging model is shown in Eq.
[Equation]
where, 𝐴𝐴𝑇𝑇 ( 𝑥𝑥 ) is ambient light based on the light attenuation of different wavelength.
𝛼𝛼 is the scene scattering coefficient, which corresponds to the scattering coefficient in the atmospheric imaging model, and 𝛼𝛼 is set by default to 1, corresponding to a moderate and homogeneous haze effect.
Three types of realistic underwater images were synthesized with color cast and haze effect are shown in Figure 1. 2.2 UWGAN FOR GENERATING REALISTIC UNDERWATER IMAGES Underwater-style images are generated based on Eq. 2, whose parameters are estimated through adversarial learning using GAN, as shown in Figure 2.
adversarial training.
2.3 U NDERWATER IMAGE RESTORATION BASED ON U-N ET
U-Net is used for color restoration and haze removal of underwater images.
A detailed description of U-Net architecture proposed in the paper is shown in Figure 3.
Firstly, a degraded underwater RGB image is resized to 256x256 and then fed into the encoder part of U-net.
In the encoder, the image is finally downsampled into a 32x32x256-dimensional latent vector through a series of convolution and max-pooling operations.
In each downsampling stage, 3x3 convolution with a stride of 1 followed by a rectified linear unit (ReLU) activation function are conducted twice, then a 2x2 max pooling with a stride of 2 is used.
The number of feature maps are doubled at each stage.
In the decoder part, upsampling is done from the latent high dimensional vector back to the original input size sequentially.
After each upsampling operation, output tensor is concatenated to the corresponding symmetric layer in the encoder side, then followed by two consecutive convolution layers and a rectified linear activation layer.
The number of feature maps is gradually reduced to three channels.
[Figure]
2.4 D ATASET
The in-air datasets we used are images of indoor scenes that has been labeled in the NYU Depth dataset V (Silberman N, 2011) and V2 (Silberman N, 2012), which contain a total of 3733 RGB images and corresponding depth maps.
The underwater dataset contains real-world underwater images collected from marine organisms’ farms (including scallops, sea cucumbers, sea urchins, etc.), which can be roughly divided into two categories, one contains near-field green hued images (RealA), and the other contains blue-green hued images of far-field scenes (RealB).
We also use underwater open datasets (Li C, 2019) (RealC) as testing sets, where RealA contains 2069 underwater images, RealB contains 2173 underwater images, and RealC contains 890 underwater images. Several typical images of the datasets are shown in Figure 4.  E XPERIMENTAL S ETUP
The training settings of our proposed method are presented in details in this section.
Our models are trained in the computer with the following configurations: Intel i7 HQ 8700 processor 16GB ， NVIDIA TITAN
X 12GB graphics card.
Firstly, UWGAN is trained to synthesize underwater-style images using the NYU-Depth Dataset, RealA and
RealB datasets.
Our model was trained for 30 epochs, using Adam optimizer with a learning rate of 0.0001, and the momentum term was set to 0.5.
The batch size was set to 64 with output images set to 256x256.
Secondly, U-net is trained as an image enhancement network using synthetic pairs.
The batch size was set to 32 and the output image size is 256x256.
The learning rate is set to 0.0001 according to Adam optimizer, our model is trained for 200 epochs.
R ESULT AND D ISCUSSION
[Figure] C o m pu t e L o ss In this section, we quantitatively and qualitatively compare our proposed method with several representative underwater image enhancement algorithms, including Unsupervised Color Correction Method (UCM) (Iqbal
K, 2010), Histogram equalization (HE) (Hummel R, 1975), Multi-Scale Retinex with Color Restoration (MSRCR) (Rahman Z, 1996), Fusion (Ancuti C, 2012), Underwater Dark Channel Prior (UDCP) (Drews
P, 2013), Image Blurriness and Light Absorption (IBLA) (Peng Y T, 2017), Underwater Color Correction using GAN (UGAN) (Fabbri C, 2018), WaterGAN-color-correction (WaterGAN) (Li J, 2017).
We employ a non-reference metric, UIQM
(Panetta K, 2015), for the quantitative assessment of underwater image quality on RealA, RealB, and RealC datasets as no ground truth scenes are available as the reference for real-world underwater images.
Besides, we employ three full-reference metrics, namely MSE, PSNR (Hore A, 2010), SSIM, for assessment image quality on synthetic datasets.
To reasonably assess the time spent on various algorithms, we resize all images to 256x256, which provides a stable output for enhancements in later experiments.
Firstly, we compare the capabilities of different methods to improve the image visibility on the RealA, RealB, and RealC datasets.
The qualitative comparison is shown in Figure 5 and 6.
Most methods can improve the quality of images of a slight haze effect.
UCM, HE, and Fusion can enhance the brightness and contrast of the image, but are less uniform for color restoration and seem to be over-enhanced in some areas of the image.
The results of MSRCR appear to have a suitable hue but lack sufficient saturation and contrast.
UDCP and IBLA do not recover well for green-toned images, they make the image darker but enhance the contrast of the image.
UGAN, WaterGAN can enhance the contrast of the image but they don’t recover color well and generate some artifacts, which destroy the structural information of the image.
The proposed method can recover the color of degraded underwater images while keeping a proper brightness and contrast.
Table 1 and Table 2 quantitatively show the scores of sample images in Figure 5 and Figure 6 respectively.
Our proposed method has achieved the highest scores in (a), (c) and (f).
In addition, the average quantized scores evaluated on RealA, RealB, and RealC datasets are shown in Table 3.
Our model achieves the best scores in terms of color restoration. Assessments
[Table] [Table]
[Table]
UIQM is a non-reference assessment metric whose quantitative results depend largely on the value of scale factors.
Structural information of images is not considered in these kinds of non-reference evaluation metrics.
Although some enhanced images can get higher score, the visual quality is poor, the reason is that the metric is calculated from the pixels.
Therefore, we also employ three full-reference assessment metrics MSE, PSNR, and SSIM to evaluate the performance of different methods on synthetic datasets without training.
The comparison results in Table 4 demonstrate that our proposed method achieves the best results in terms of MSE, PSNR, and SSIM.
The average inference time of different algorithms are compared in one computer with following configuration:
Intel i7-8750H CPU, 16GB RAM, and GTX1060 6G GPU.
The results are shown in Table 5.
Our model has the fastest processing speed compared to other methods.
Moreover, the model we proposed has the fewest Params and FLOPs compared to other deep-learning-based methods.
UGAN employs many convolution layers with 512 kernels, which causes that there are too many network parameters.
WaterGAN employs multiple networks, resulting in slow processing speed.
[Table] [Table]
indicated by some previous works (Uplavikar P, 2019; Anwar S, 2019; Ding X, 2019), the performance of high-level computer vision tasks (such as underwater target detection) on enhanced images is an indicator of image enhancement methods. We applied YOLO v3 (Redmon J, 2018) target detector on degraded underwater images and their enhanced versions generated by our model.
The performance of underwater target detection is better on enhanced versions on degraded images. Figure 7 shows the results of YOLO v detector before and after processing the images with our model.
C ONCLUSION
Based on an improved underwater imaging model, a generative adversarial network (UWGAN) for generating realistic underwater images is proposed in this paper.
Then, U-net with combined loss functions is used for degraded underwater images enhancement.
Our model is validated on both low-level and highlevel underwater computer vision tasks, which demonstrate its effectiveness and robustness.
A CKNOWLEDGEMENT
The authors would like to acknowledge the National Key R&D Program of China (Grant No.
2018YFC0309402) for funding this work.
A L OSS F UNCTIONS
The most common loss function for image restoration is L2 error.
However, which loss function is suitable for underwater image enhancement has not been studied.
Inspired by a related article, the effect of different loss functions in U-net is studied in this paper.
Table 6 shows the loss functions we used.
In mathematical formula, 𝑥𝑥 is an index of pixels in region 𝑋𝑋 , 𝑔𝑔 ( 𝑥𝑥 ) is pixel value in region 𝑋𝑋 of the image reconstructed by U-net and 𝑟𝑟 ( 𝑥𝑥 ) is the pixel value of corresponding ground truth.
𝑥𝑥 is the central pixel value of region 𝑋𝑋 .
∇𝑔𝑔 ( 𝑥𝑥 ) , ∇𝑟𝑟 ( 𝑥𝑥 ) respectively represent the gradient of reconstructed images and clear images.
After several experiments and observations of the best reconstruction results, we set 𝛼𝛼 to 0.8 in this paper.
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation] [Figure]
(i) are respectively enhanced results of the loss function
[Equation]
𝑆𝑆𝑆𝑆𝐼𝐼𝑆𝑆 , 𝐿𝐿 1 + 𝑆𝑆𝑆𝑆𝑆𝑆𝐼𝐼𝑆𝑆 , and 𝐿𝐿 1 + 𝐺𝐺𝐺𝐺𝐿𝐿 .
Ablation study is mainly to reveal the effects of different loss functions.
We use different loss functions to train the network and test it on RealC, and synthetic datasets.
The sample image is selected from the RealC dataset, as shown in Figure 8, the enhanced results range from (a)~(i) are obtained with different loss functions, and images in the second row show the details in the red box area of the image.
It can be seen from the results in the second row, (a), (b), (f) appear striped artifacts.
(c), (d), (g) cause color unevenness.
The details of (e) are natural but it lacks sufficient saturation.
(h), (i) show proper enhanced results, the color in (h) is more vivid but with slightly striped artifacts.
The enhanced result using the 𝐿𝐿or 𝐿𝐿loss function appears stripe-like artifacts while the 𝑆𝑆𝑆𝑆𝐼𝐼𝑆𝑆 or 𝑆𝑆𝑆𝑆𝑆𝑆𝑆𝑆𝐼𝐼𝑆𝑆 loss function causes color unevenness.
The enhanced result of the 𝐺𝐺𝐺𝐺𝐿𝐿 loss function is natural but lacks sufficient saturation.
We calculated the MSE, PSNR, and SSIM metrics on the synthetic dataset.
The quantitative scores in Table 7 demonstrate that a combination of multiple loss functions can achieve better enhancement results.
[Table]