UCT‐GAN: underwater image colour transfer generative adversarial network

IET Image ProcessingVolume 14, Issue 14 p. 3613-3622 Research ArticleFree Access UCT-GAN: underwater image colour transfer generative adversarial network Junjie Deng, School of Informatics Xiamen University, Xiamen, 361005 People's Republic of ChinaSearch for more papers by this authorGege Luo, School of Electronic Science and Engineering, Xiamen University, Xiamen, 361005 People's Republic of ChinaSearch for more papers by this authorCaidan Zhao, Corresponding Author zcd@xmu.edu.cn School of Informatics Xiamen University, Xiamen, 361005 People's Republic of ChinaSearch for more papers by this author Junjie Deng, School of Informatics Xiamen University, Xiamen, 361005 People's Republic of ChinaSearch for more papers by this authorGege Luo, School of Electronic Science and Engineering, Xiamen University, Xiamen, 361005 People's Republic of ChinaSearch for more papers by this authorCaidan Zhao, Corresponding Author zcd@xmu.edu.cn School of Informatics Xiamen University, Xiamen, 361005 People's Republic of ChinaSearch for more papers by this author First published: 12 November 2020 https://doi.org/10.1049/iet-ipr.2020.0003AboutSectionsPDF ToolsRequest permissionExport citationAdd to favoritesTrack citation ShareShare Give accessShare full text accessShare full-text accessPlease review our Terms and Conditions of Use and check box below to share full-text version of article.I have read and accept the Wiley Online Library Terms and Conditions of UseShareable LinkUse the link below to share a full-text version of this article with your friends and colleagues. Learn more.Copy URL Share a linkShare onEmailFacebookTwitterLinked InRedditWechat Abstract Underwater image enhancement algorithms improve image quality and indirectly enhance underwater visibility. Although many underwater image enhancement neural networks have been proposed, they require large amounts of data. To reduce the amount of data required while providing better image enhancement, this study proposes an underwater image colour transfer generative adversarial network (UCT-GAN). The authors first design a non-linear mapping function to generate colour cast images according to original images. Then, the authors utilise these image pairs (i.e. colour cast images and corresponding original images) to guide the UCT-GAN in learning the inverse function of the designed non-linear mapping function. Finally, colour cast images are restored via the inverse function. A data augmentation method based on Poisson fusion and block combination is also proposed to overcome the problem of requiring a large amount of training data. Moreover, the authors extend UCT-GAN into a multi-class colour transfer network to achieve an array of underwater image enhancements. Experimental results indicate that the proposed UCT-GAN can more effectively resolve underwater image colour cast compared to existing algorithms. 1 Introduction When people acquire pictures through different image acquisition devices, they need to adjust the acquisition device parameters, such as white balance and gamma correction parameters, according to changes in the environment. Adjusting the settings of image acquisition equipment requires professional skills and rich experience, making it impossible for general users to obtain high-quality pictures efficiently and quickly. Image enhancement algorithms enable users to get swiftly satisfactory enhanced pictures, which can be divided into four categories: methods based on histogram equalisation [1, 2], methods based on Retinex theory [3–6], algorithms based on dark channel priors [7–9] and deep learning [10, 11]. Many methods have recently been proposed to achieve underwater picture enhancement. The histogram equalisation algorithm is the most basic linear mapping method. Hitam et al. [2] proposed a mixed contrast limited adaptive histogram equalisation colour model, which was found to substantially improve the visual quality of underwater images. Mathur and Goel [3] used the automatic white balance algorithm to stretch the red, green and blue three-channel histogram of a picture to different degrees and introduced the gamma correction algorithm to improve image contrast, which resolved the colour cast problem to some extent. However, these statistical methods often suffer from excessive or insufficient enhancements in underwater image correction, which introduces artefacts and amplifies noise in images in practical applications. Fu et al. [4] introduced a basic grey world colour correction method into underwater image processing, which prevented images from being over-enhanced by the control coefficient ; they also used the fast Fourier transform algorithm to accelerate the processing speed of underwater image enhancement. Li et al. [5] applied the dark channel dehazing [6] algorithm to underwater image enhancement and solved the problem of underwater image colour cast via colour-correction processing based on the grey-world hypothesis. However, due to low contrast and less visible edges in underwater images, assumptions of the grey-edge hypothesis algorithm are destroyed in cases of insufficient illumination, thus rendering image enhancement unsatisfactory. Applying the dark channel dehazing algorithm to underwater image enhancement requires accurate estimation of background light. Wen et al. [7] developed a new underwater optical model to describe the process of underwater imaging and proposed an effective image enhancement algorithm that can handle images taken in the deep sea and turbid waters. Bokar and Bonde [8] put forth a red-channel underwater enhancement algorithm based on dark channel theory and established a depth map using morphological methods to realise underwater image enhancement. Convolutional neural networks (CNNs) have achieved great success in the modelling of complex non-linear systems. They have also demonstrated satisfactory results in image processing, but the volume of CNN parameters is vast and requires a large amount of training data for a better model fit [9]. To mitigate this problem, Sun et al. [10] continued to train a pre-training model with an underwater image set to accomplish underwater image enhancement according to the transfer learning method. Wang et al. [11] synthesised underwater images on the basis of the underwater image imaging model [12] to minimise insufficient data in the CNN training process. Li et al.'s [13] WaterGAN algorithm uses the noise vector, pictures taken in air, and corresponding picture depth as input; the trained WaterGAN can transform pictures taken in the air into underwater pictures to realise end-to-end learning. Also, UWGAN [14] is an improved version of the WaterGAN, and this method can enhance underwater images effectively. Chen et al. [15] built the generative adversarial network (GAN)-RS to restore the distorted images in terms of dark channel prior. Their network can adaptively improve underwater visual quality in real-time and induce an overall superior restoration performance. Overall, traditional underwater image enhancement methods have numerous limitations in application scenarios and experimental conditions, and the deep learning method requires a large amount of training data. To overcome these drawbacks, this paper introduces an underwater image enhancement method based on colour transfer. A multi-class colour transfer enhancement network that uses global and local colour information is designed to realise few-shot learning; resultant enhanced images exhibit improvements in colour, contrast and naturalness. To better accomplish multi-class colour transfer, this paper takes large margin cosine loss (LMCL) [16] as a classification metric. Poisson fusion [17] and block combination are introduced to realise few-shot learning, which can better preserve detailed image contents. Additionally, this paper constructs a brightness discriminator to help eliminate unevenness in target-image brightness, further improving enhanced picture quality. The main contributions of this paper are as follows: • We build a novel GAN architecture which can implement multi-class colour transfer. • We propose a kind of data augmentation methods that realise few-shot learning. • We design a light discriminator (LiD) which improve enhanced image quality. The remainder of this paper is organised as follows. Work related to the proposed algorithm is introduced in Section 2. Section 3 presents detailed algorithm formulas, principle flow diagrams and data augmentation methods; network structure, training setup and training loss are outlined in Section 4. Section 5 provides an experimental comparison and data analysis, and Section 6 summarises our work and identifies potential avenues for future research. 2 Related work GANs have achieved satisfactory results in areas such as image generation [18], super-resolution [19], image inpainting [20] and style transfer [21]. A typical GAN consists of two parts: a generator and a discriminator. The generator generates a false sample based on the feature distribution of the real sample, and the discriminator judges whether the input sample is a real or false sample generated by the generator. More specifically, the GAN implements corresponding tasks according to sample characteristics. In one case, the features of generative pictures should be as similar as possible to the real sample to fool the discriminator, thus training a more powerful generator. The reverse is true in the other case, intended to obtain a better discriminator. In the field of image generation, input refers to noise and output refers to an image; as such, the absence of an image-to-image training method requires the GAN to learn features from a large number of real samples. Image generation is therefore challenging. In areas such as style transfer, GAN implements an image-to-image task, which facilitates simple problem resolution. An image-to-image network is widely used in the field of style transfer, wherein the generator realises inter-domain mapping to complete the transformation of the picture style. Underwater picture distortion is complicated; obtained pictures may be blue as well as red, yellow and dark. The style transfer method can transform images with different distortions into the same style type, effectively addressing diverse colour distortion problems. Li et al. [22] introduced the SSIM [23] loss metric based on CycleGAN [24] and achieved colour enhancement of underwater images by limiting the texture edges of generated images. Their method was simple and effective, and it demonstrates the feasibility of using GAN for colour transfer to complete underwater image enhancement tasks. Fabbri et al. [25] used CycleGAN to convert pictures taken in air into underwater images. This method builds a paired data set that implements the end-to-end learning approach via UGAN. End-to-end learning uses the GAN network to correct underwater images based on the original image. Fig. 1 suggests that the advantages and disadvantages of end-to-end learning. End-to-end learning (also called end-to-end training) is the method which directly optimise the system target without training in modules or phases during the learning process. Its training data follows the form of 'input–output' pairs and the error obtained from comparing the real result with the output will be transmitted at each layer in the model to adjust the network parameters. This process will not end until the model converges or achieves the expected effect. Although end-to-end learning is simple, a lot of massive labelled data is essential. To generate massive data, Li et al. [13] and Wang et al. [14] build special generators to synthesise underwater images based on in-air images respectively. Inspired by their work, the underwater image colour transfer GAN (UCT-GAN) also adopts the end-to-end learning and we proposed a new method to generate labelled data. Fig. 1Open in figure viewerPowerPoint Two learning approaches 3 Proposed method To realise the end-to-end learning process, we propose a simple and effective method for acquiring underwater colour distortion pictures. Using this method, we can simulate the underwater image colour cast and construct GAN to learn its inverse process. To surmount the obstacle of insufficient training data in real-life scenarios, we apply an image fusion algorithm and a picture block combination method to further improve model performance. 3.1 Solution of colour correction In image enhancement, the generalised linear monotonically increasing mapping function does not change the size relationship between original pixels. A simple but effective image enhancement method is available to complete colour correction by using a function that satisfies the above property. The formula can be expressed as: (1)where and represent the colour-enhanced picture and cast picture, respectively; and is a generalised linear monotonically increasing mapping function whose parameters are defined as . When we fix the position of the image acquisition device to obtain a picture with different colour temperatures in the same scene and calculate the corresponding colour mapping functions, we find that the colour cast mapping function increases non-linearly because the function is not a continuous line. The obtained mapping function is shown in Fig. 2 Using the linear least squares method to fit these discrete points, the expression can be written as: (2)where represents the original image, and is a solvable generalised linear monotonically increasing mapping function. The result obtained using the linear least squares algorithm does not precisely represent true discrete points; rather, it approximates the distribution trend of these points. To more effectively correct the colour-distorted image, we aim to solve a non-linear mapping function that can realise non-linear mapping based on the linear least squares fitting function. The new formula is similar to (2) (3)where M is a set of random distributions that minimises (3). We want to add a random variable based on the generalised linear monotonically increasing function to simulate the actual mapping function, which can transform the mapping function into a non-linear form, thus achieving better colour correction. The new problem is that, although there is a non-linear mapping function including a specific distribution m and a linear function to minimise (3), a closed solution remains difficult to identify. Another problem is that fewer pairs of images consisting of original and unbalanced pictures make end-to-end learning impossible. Since the imbalance is relative, according to (3), the distorted picture can be corrected to the original picture, and the original picture can be converted into a colour-distorted picture. Therefore, the following colour imbalance mapping function is designed to simulate the underwater picture: (4)where is either an affine function, a logarithmic function or an exponential function. These generalized linear monotonically increasing functions can preserve the original colour trend relative to size. In this case, m is a Gaussian-distributed random variable used to simulate non-linear mapping. is a non-linear function whose parameters are defined as , containing p unknown parameters. The parameters of the function and n are random; in this paper, they all satisfy the uniform distribution. Therefore, our goal is to obtain a paired image through the constructed distortion function and then use its inverse function to restore the original image from the distorted image and realize end-to-end learning. Finally, the inverse function will restore the original image from the distorted image. According to (3) and (4), the final solution formula is as follows: (5) Fig. 2Open in figure viewerPowerPoint Mapping function between pictures with different colour temperatures From (5), we see that the parameters of and m are unknown, and solving the inverse function is challenging. Therefore, we use a CNN to approximate the non-linear function. Its multiple hidden-layer neuron connections cause the entire network to demonstrate good fitting performance. CNN's have been shown to fit arbitrary non-linear functions. Thus, using CNN suggests that applying the inverse function of non-linear mapping is feasible. Further, GAN is used in this study to implement the inverse function learning process described above. Fig. 3 presents a flowchart of the proposed algorithm, which generates distorted images through the colour cast mapping function in (4). Then, the GAN can learn the inverse function of the colour cast function to complete the colour correction process for the distorted picture. The optimisation goal is to minimise the mean square variance (MSE) of original and enhanced images as much aspossible. Fig. 3Open in figure viewerPowerPoint Algorithm flow chart 3.2 Data augmentation We can generate an infinite number of image pairs composed of original images and corresponding distorted images according to (4), but when only a single original image is used to obtain a large amount of training data, the test results are not ideal. Therefore, we introduce the idea of block combination to increase training picture textures in an attempt to improve the quality of generated pictures. This method first randomly crops the image by 32 32 size and then re-splices the randomly flipped cropped images into a new 256 256 resolution image, which increases the texture of each image in seven directions. However, experimental results show that the generated pictures are slightly distorted and blurred when the amount of original data is small. Thus, the Poisson fusion [17] method is introduced to solve this problem. Fig. 4 displays the data augmentation flow chart. The algorithm selects the region to be fused by the mask and solves the Poisson equation of the fusion region, after which target images are embedded according to colour information of the fusion region. The Poisson fusion algorithm can preserve the colour information of the original image while introducing texture features without incorporating excessive colour-related information. Fig. 5 illustrates the experimental results of the data augmentation method proposed herein. Through an ablation experiment, we found that only using a large number of pictures collected at the same colour temperature for training could complete the image enhancement task; however, some details of the corrected images were missed. By using the proposed method to train the model, image details improved because the fused image contained more texture-related information. Colour distortion is global, and block combination operations neither affect global mapping nor produce a block effect in the generated image; instead, they help the model learn more subtle features from each tile. Fig. 4Open in figure viewerPowerPoint Data augmentation flow chart Fig. 5Open in figure viewerPowerPoint Experimental results of data augmentation (a) Original image, (b) Test results of images without fusion or block combination (based on one training image), (c) Test results of images without block combination (based on one training image with fusion), (d) Test results of images without fusion (based on one training image with block combination), (e) Test results of images (based on one training image with fusion and block combination), (f) Test results of images without fusion or block combination (based on many training images) 4 Underwater image colour transfer GAN In this section, the detailed UCT-GAN architecture and training setup are proposed. The adversarial and generative loss is also introduced. 4.1 Generative and adversarial loss We obtain paired training images through (4) and then implement the end-to-end learning process. The MSE loss is introduced to measure the colour correction effect. The loss expression is as follows: (6) In (6), K represents the number of image categories. Training the network through paired images fits the network parameters well and guides neural network training; however, the initial experimental results revealed that the brightness information was not transferred well, resulting in darker generated images. Therefore, LiD was introduced to measure brightness loss and promote effective brightness transfer. Moreover, because the proposed algorithm trained the model based on very little data, the training picture was lacking and the network could not acquire statistical brightness information. To solve this problem, UCT-GAN first calculates the average brightness of each original image category and then discriminates the average brightness of generated images. Accordingly, an image with the average brightness difference of the training picture within a certain size can be transferred successfully; otherwise, the brightness transfer will fail. The brightness loss expression is as follows: (7) In (7), the random variable z is used to limit the brightness range. We also construct a multi-colour transfer network to address the diversity of choices. Combining the dimension-expanded label with the image enables UCT-GAN to perform multiple types of colour transfer enhancements on an image based on combined label information. We also used the LMCL to calculate classification loss. The LMCL algorithm converts a traditional softmax loss to a cosine loss by (i.e. by normalising features and weight vectors to remove radial variations) and then introduces a cosine margin term to further maximise the decision spacing of different classes in the angular space. The expression is as follows: (8) In (8), N is the number of training samples, equal to 1 in this paper. is the ith eigenvector corresponding to the real category ; is the weight vector of the jth category and is the angle between and . We also use global discriminator (GD) to calculate the overall transfer performance of the image colour and use LiD to calculate the colour transfer performance of partial images. The two discriminators each use WGAN-GP to calculate the loss, the expression of which can be written as follows: (9) (10) In (9) and (10), is the gradient operator; denotes the local discriminator (LOD); denotes the GD; and is the weight coefficient, which is equal to 10 in this paper. The LiD can discriminate the colour cast of the image detail, enabling better colour-related transfer for that portion. The GD can discriminate colour information within the whole picture to ensure comprehensive colour transfer. The total discriminator and generator loss can be written as: (11) (12) In (11) and (12), represents the loss of ground truth distribution, and represents the classification loss of generated false images. Essentially, the optimisation goal of the generator is to minimise (11) and that of the discriminator is to minimise (12). 4.2 Training detail We use DenseNets [26] modules as the main structure. DenseNets can effectively transmit feature information, reduce the number of parameters, and reduce gradient disappearance. Using the deconvolutional layer as the up-sampling operation results in a checkerboard effect and deteriorates the quality of the generated image. Therefore, we combined the up-sampling layer and the convolutional block to realise up-sampling. Each convolutional block architecture uses the same combination of a convolutional layer and switchable normalisation [27] followed by the leaky rectified linear units activation layer. SN is the improved combination of batch normalisation, instance normalisation and layer normalisation, and it can learn to select different normalisers for different normalisation layers of a deep neural network. SN computes statistics (means and variances) from different dimensions, such as minibatch, layer and channel. At last, SN switches the three normalisers by learning their weights. The Adam optimiser includes and , the learning rate is fixed at 0.002, and the batch size is fixed at 1. The intensity of the colour cast function increases as the training epoch increases, which can alleviate the model's oscillation intensity when training begins. We also use WGAN-GP [28] in local and GDs to stabilise the model and ensure that the model parameters can converge normally and improve the quality of generated images. WGAN-GP is an improved algorithm from WGAN and a variant of GAN, which contributes to solving the problem of WGAN [29] on training difficulty and slow convergence in the actual experimental process. Specifically, in connection with the matter of direct adoption of weight clipping when dealing with the Lipschitz constraints leading to gradient disappears or gradient explodes, WGAN-GP removes it and adds gradient penalty to discriminator loss. In the training stage, the system randomly selects one original image that is from some class, then the Poisson fusion and block combination algorithms will be implemented to generate ground truth. Finally, we can get massive distorted images as inputs in terms of the (4). In the test stage, we select some class as enhancement direction. Thus, all kinds of colour-transfer images can be generated. Fig. 6 shows the UCT-GAN presented in this paper. In terms of the generation model, we first use two convolutional blocks to reduce the image resolution and increase its depth. Then, the DenseNets modules are used to extract image features. Finally, two up-sampling blocks restore the original image size. The discrimination model is divided into a LoD, GD and LiD. The LoD and GD respectively determine local and global colour information of the image to achieve better colour transfer enhancement; the LiD ensures that the image brightness information is transferred more effectively. Fig. 7 shows the comparison between the proposed model and others' models. UWGAN and WaterGAN both can generate underwater images by training special GANs. Then, an end-to-end CNN can be trained according to these generated images. Finally, the image enhancement task can be completed in the test step. Compared with their methods, our algorithm can complete the image augmentation task by one step, which shows that our method is simpler. Besides, UCT-GAN demonstrates a good multi-class colour transfer enhancement effect and UWGAN or WaterGAN only realises single-class image enhancement. Fig. 6Open in figure viewerPowerPoint UCT-GAN architecture (a) LiD, (b) GD, (c) LoD Fig. 7Open in figure viewerPowerPoint Model comparison 5 Performance analysis In this section, we present training and test data used in this paper along with the UCT-GAN underwater image enhancement results using the brightness discriminator. We found that the constructed brightness discriminator can effectively improve the brightness and quality of pictures through an ablation experiment. Meanwhile, We provide the non-reference image quality evaluation (NIQE) scores and test results of UCT-GAN. Compared with various algorithms, UCT-GAN is found to exhibit effective underwater image enhancement. In Section 5.3, We give more ablation tests to support the effectiveness of the proposed method. 5.1 Data set Fig. 8 illustrates image acquisition platform that consists of the TES-136 CHROMA METER, one fixed camera and a table lamp with adjustable colour temperature. We took ten indoor images with colour temperatures of 3500, 5500 and 7500 K as source data and classified the original pictures into Categories 1, 3 and 4, respectively; five pictures with an average brightness of 60 were randomly selected as the source data for Category 2. For convenience, 'TF1', 'TF2', 'TF3' and 'TF4', respectively, denote the four types of enhanced pictures and 'TFn+light' indicates the nth-class images enhanced by the brightness discriminator. We used roughly 17,000 images from the voc2012 database as a fused image for texture augmentation. The proposed algorithm can arbitrarily select a small number of pictures with similar colour temperatures as domain training data. The experiment included three large data sets for testing: 1956 underwater images on the Flickr website; 15,000 underwater images from the University of Michigan's Marine Fluid Dynamics Laboratory (MHL) [30] and 18,000 images from Port Royal, Jamaica, at the site of a submerged city containing natural and man-made structures [30]. Fig. 8Open in figure viewerPowerPoint Image acquisition platform 5.2 Experimental results and analysis Fig. 9 indicates that after using the brightness discriminator, details of the generated picture became more prominent, and the picture contrast was improved to some extent. These outcomes were presumably due to effective transfer of brightness information for the detailed image parts; therefore, the brightness of details did not produce a blurred effect and thus enhanced the quality of the colour-transfer-enhanced picture. Table 1 shows the average brightness of different categories under the three test data sets, where '+light' indicates cases in which the LiD was used. After using LiD, the average brightness of generated pictures increased. Fig. 10 displays the partial underwater image enhancement effect and each column from left to right is original image, 'cyclegan-ssim', 'Fu', 'UGAN', 'WaterGAN', 'GAN-RS', 'UWGAN', 'TF1+light', 'TF2+light', 'TF3+light' and 'TF4+light'. UCT-GAN could generate enhanced images with different colour temperatures and brightness levels according to the input label; our proposed network could also enhance the image texture to be more detailed with higher contrast and better colour correction. Fig. 9Open in figure viewerPowerPoint Colour enhancement contrast effect afte