All-In-One Underwater Image Enhancement using Domain-Adversarial Learning

Underwater images have an application in a variety of fields like marine research and underwater robotics. We need clear underwater imagery to study deteriorating coral reefs and other aquatic life. Underwater robotic systems also rely heavily on high quality images to fulfill their objectives. However, the quality of the images acquired for these applications is degraded due to various factors. One of the major factors for this degradation is wavelength dependent light attenuation over the depth of the object in the scene. For example, red light is absorbed in water at a higher rate than blue or green light. Hence, we see a blueish or a greenish tint in an underwater scene. Another factor diminishing underwater image quality is the light scattered due to the small particles present in water, which introduces a homogeneous background noise to the image.

Apart from these factors, another challenge in underwater image enhancement is the diversity of underwater image distributions. We can see this diversity in figure 1, which shows how underwater scenes captured in shallow coastal waters look different than those captured in deep oceanic waters or those captured in muddy waters. It is hard for a single model to enhance underwater images for such multiple image distributions and, therefore, providing a universal solution for underwater image enhancement is difficult. While previous work has addressed the challenges of light attenuation and scattering, not many have handled the challenge of image distribution diversity explicitly. [2] proposes color restoration of underwater images by performing color correction using attenuation coefficient ratios for all the Jerlov water types and then selecting the best result out of them. Whereas, [3] proposes one solution by training multiple models, each for a different Jerlov water type. But these approaches seem inefficient and rely on the prior knowledge of the water type for the given image to perform color restoration.

One more challenge faced in underwater image enhancement is the lack of real-world datasets containing the ground truth clear images, as it is extremely difficult to find degraded and clear versions of the same real-world underwater scene, thus creating a bottleneck for data-driven methods. To address this challenge, synthetic underwater image datasets have been built in the past. One such example is the work done in [3], who synthesize an underwater image dataset consisting images of 10 Jerlov water types using the NYU Depth Dataset V2 [7] which provides the ground truth clear images and the depth, both of which are required to synthesize degraded underwater images using the underwater image formation model [8].

The task of enhancing underwater images is, therefore, difficult and has its own unique challenges. These images fail at many vision tasks like object detection, classification, segmentation which provokes a need to process them and enhance their quality. We propose a novel solution to this problem by addressing all the challenges mentioned above using a convolutional neural network [9] based encoder-decoder to reconstruct clear underwater images and a convolutional neural network based classifier to classify the Jerlov water types, which acts as our nuisance classifier. We synthesize a dataset of synthetic underwater images by following [3] to train our model, thus taking into account the factors of wavelength dependent light attenuation and scattering of light due to small water particles while forming an underwater image and thereby addressing the respective challenges. To address the challenge of underwater image distribution diversity, we train our model to learn the domain agnostic features for a given degraded underwater image, where the domain is the Jerlov water type of the image. The objective of our encoder, apart from learning an encoding to reconstruct a clear underwater image, is to make the prediction of the nuisance classifier as uncertain as possible by discarding the features denoting the water type and preserves only the scene related features [10]. This is similar to a generator fooling a discriminator in a generative adversarial network [11], the only difference being that in this case we want to make the classifier uncertain. We, therefore, introduce an adversarial loss on our encoder, computed at the end of the nuisance classifier, which is negative entropy. The encoder is also acted upon by a reconstruction loss computed at the end of our decoder. Thus, we propose a novel model which performs underwater color restoration for multiple types of underwater images.

Our approach has multiple highlights: i) Our proposed approach is able to learn water-type agnostic features. We adapt the adversarial training strategy proposed in [10] to our model; ii) Our proposed model outperforms the previous enhancement methods in SSIM and PSNR scores for almost all Jerlov water types; iii) Our model has good generalization ability on real-world datasets, as well as performs nicely on improving subsequent object detection on enhanced images.

Many previous attempts to solve the underwater image enhancement problem have used physics-based methods. [12] tries to solve this problem by explicitly modeling the refraction in water, whereas, [13] incorporates the inherent properties of the underwater medium such as attenuation, scattering, and the volume scattering function in order to simulate image formation. [8] defines an underwater image formation model which is given as

where $U_c\left(x\right)$ is a point $x$ in the underwater image, $I_c\left(x\right)$ is a point $x$ in the clear image, $T_c\left(x\right)$ is the fraction of the light reaching the camera after reflecting from point $x$ in the scene and $B_c$ is the homogeneous background light of the scene. $T_c\left(x\right)$ is further given as

where ${\beta }_c$ is the wavelength dependent medium attenuation coefficient, $E_c(x,d(x\left)\right)$ is the energy of a light beam from point $x$ after it passes through a medium and $N_c\left(d\right(x\left)\right)$ is the normalized residual energy ratio for every unit of depth covered.

The above physical model is similar to that of image dehazing, except that the medium attenuation coefficient is wavelength dependent, whereas in dehazing it does not depend on the light wavelength. This model has been used by many approaches to solve the underwater image enhancement problem. [14] tries to improve on the above model by computing attenuation coefficients in the 3D RGB space, whereas [3] uses the above model to generate a synthetic dataset of 10 Jerlov water types. We generate a similar dataset in our work, the details of which are given in section 4.2.1.

In recent years, deep learning [15] techniques like Convolutional Neural Networks (CNN) [9] and Generative Adversarial Networks (GAN) [11] have been very effective at solving vision problems. Naturally, these techniques have then been used for underwater image enhancement. [16] trains a GAN to learn the mapping from underwater to clear images. [3] train multiple CNN models, each for different water type in their dataset, to get enhanced images. However, these methods fail to provide a singular solution capable of handling the diversity of underwater images apart from generating their clear versions.

Since one of our goals, apart from underwater image enhancement, is to train a single model which can do this task for multiple water types, we first try to learn a water type agnostic encoding for the given underwater image. That means, ideally, the latent vector $Z$ extracted from an encoder $E$ for the same underwater scene, should be the same for different water types. That way the decoder or the generator $G$ is able to reconstruct a clear image of the scene from only the scene specific features. Both $E$ and $G$ are neural networks in our model.

To do so, we introduce a novel application of a nuisance classifier $D$ along with $E$ and $G$. The nuisance classifier is a neural network which aims to classify the water type of the given input image from its latent vector $Z$ extracted from the encoder. However, we also introduce an adversarial loss [11] over the encoder using the nuisance classifier. Our formulation of the adversarial loss forces the encoder to generate $Z$ such that the nuisance classifier is unsure of the possible water types. Thus, the adversarial loss forces the encoding to be agnostic of the features denoting the water type. The full architecture can be seen in figure 2.

We train our model on the synthetic underwater image dataset described in detail in section 4.2.1. The model is trained on a machine with the following configuration - Intel i7 6700 HQ processor, 8 GB RAM, NVIDIA GeForce GTX 960M 4GB graphics card.

We are able to provide a novel solution for underwater image enhancement which outperforms the previous methods both qualitatively and quantitatively. Our goal is to provide a generalized solution which could handle the diversity of the underwater images as well as transform them into clear images. Our model is successful in doing so by learning domain agnostic features for multiple underwater image types and then generating their clear version from those features. We also show that the model is able to generalize well on the unseen real-world data. Also, experimental results on object detection task show that enhancing underwater images with our model before high level vision tasks improves the detection performance.