Maintaining Natural Image Statistics with the Contextual Loss

Introduction

“Facts are stubborn things, but statistics are pliable.” ― Mark Twain

Maintaining natural image statistics has been known for years as a key factor in the generation of natural looking images . With the rise of CNNs, the utilization of explicit image priors was replaced by Generative Adversarial Networks (GAN) , where a corpus of images is used to train a network to generate images with natural characteristics, e.g., . Despite the use of GANs within many different pipelines, results still sometimes suffer from artifacts, balanced training is difficult to achieve and depends on the architecture .

Well before the CNN era, natural image statistics were obtained by utilizing priors on the likelihood of the patches of the generated images . Zoran and Weiss showed that such statistical approaches, that harness priors on patches, lead in general to restoration of more natural looking images. A similar concept is in the heart of sparse-coding where a dictionary of visual code-words is used to constrain the generated patches . The dictionary can be thought as a prior on the space of plausible image patches. A related approach is to constrain the generated image patches to the space of patches specific to the image to be restored . In this paper we want to build on these ideas in order to answer the following question: Can we train a CNN to generate images that exhibit natural statistics?

The approach we propose is a simple modification to the common practice. As typically done, we as well train CNNs on pairs of source-target images by minimizing an objective that measures the similarity between the generated image and the target. We extend on the common practice by proposing to use an objective that compares the feature distributions rather than just comparing the appearance. We show that by doing this the network learns to generate more natural looking images.

A key question is what makes a suitable objective for comparing distributions of features. The common divergence measures, such as Kullback-Leibler (KL), Earth-Movers-Distance and $\chi^{2}$, all require estimating the distribution of features, which is typically done via Multivariate Kernel-Density-Estimation (MKDE). Since in our case the features are very high dimensional, and since we want a differentiable loss that can be computed efficiently, MKDE is not an option. Hence, we propose instead an approximation to KL that is both simple to compute and tractable. As it turns out, our approximation coincides with the recently proposed Contextual loss , that was designed for comparing images that are not spatially aligned. Since the Contextual is actually an approximation to KL, it could be useful as an objective also in applications where the images are aligned, and the goal is to generate natural looking images.

Since all we propose is utilizing the Contextual loss during training, our approach is generic and can be adopted within many architectures and pipelines. In particular, it can be used in concert with GAN training. Hence, we chose super-resolution as a first test-case, where methods based on GANs are the current state-of-the-art. We show empirically, that using our statistical approach with GAN training outperforms previous methods, yielding images that are perceptually more realistic, while reducing the number of required training images by orders of magnitude.

Our second test-case further proves the generality of this approach to data that is not images, and shows the strength of the statistical approach without GAN. Maintaining natural internal statistics has been shown to be an important property also in the estimation of 3D surfaces . Hence, we present experiments on surface normal estimation, where the network’s input is a high-resolution image and its output is a map of normals. We successfully generate normal-maps that are more accurate than previous methods.

To summarize, the contributions we present are three-fold:

We show that the Contextual loss can be viewed as an approximation to KL divergence. This makes it suitable for reconstruction problems.

We show that training with a statistical loss yields networks that maintain the natural internal statistics of images. This approach is easy to train and reduces significantly the training set size.

We present state-of-the-art results on both perceptual super-resolution and high-resolution surface normal estimation.

Training CNNs to Match Image Distributions

Our approach is very simple, as depicted in Figure 2. To train a generator network $G$ we use pairs of source $s$ and target $y$ images. The network outputs an image $G(s)\!\equiv\!x$, that should be as similar as possible to the target $y$. To measure the similarity we extract from both $x$ and $y$ dense features $X\!=\!\{x_{i}\}$ and $Y\!=\!\{y_{j}\}$, respectively, e.g., by using pretrained network or vectorized RGB patches. We denote by $P_{Y}$ and $P_{X}$ the probability distribution functions over $Y$ and $X$, respectively. When the patch set $X$ correctly models $Y$ then the underlying probabilities $P_{Y}$ and $P_{X}$ are equal. Hence, we need to train $G$ by minimizing an objective that measures the divergence between $P_{X}$ and $P_{Y}$.

2 Review of the KL-divergence

There are many common measures for the divergence between two distributions, e.g., $\chi^{2}$, Kullback-Leibler (KL), and Earth Movers Distance (EMD). We opted for the KL-divergence in this paper.

The KL-divergence between the two densities $P_{X}$ and $P_{Y}$ is defined as:

It computes the expectation of the logarithmic difference between the probabilities $P_{X}$ and $P_{Y}$, where the expectation is taken using the probabilities $P_{X}$. We can rewrite this formula in terms of expectation:

This requires approximating the parametrized densities ${P_{X}},{P_{Y}}$ from the feature sets $X\!=\!\{x_{i}\}$ and $Y\!=\!\{y_{j}\}$. The most common method for doing this is to use Multivariate Kernel Density Estimation (MKDE), that estimates the probabilities $P_{X}(p_{k})$ and $P_{Y}(p_{k})$ as:

Here $K_{H}(p_{k},\xi)$ is an affinity measure (the kernel) between some point $p_{k}$ and the samples $\xi$, typically taken as standard multivariate normal kernel $K_{H}(z)=(2\pi)^{-d/2}|H|^{-1/2}\exp(-\frac{1}{2}z^{T}Hz)$, with $z=dist(p_{k},\xi)$, $d$ is the dimension and $H$ is a $d\!\times\!d$ bandwidth matrix. We can now write the expectation of the KL-Divergence over the MKDE with uniform sample grid as follows:

A common simplification that we adopt is to compute the MKDE not over a regular grid of $\{p_{k}\}$ but rather on the samples $\{x_{i}\}$ directly, i.e., we set $p_{k}\!=\!x_{k}$, which implies $K_{H}(p_{k},x_{i})\!=\!K_{H}(x_{k},x_{i})$ and $K_{H}(p_{k},y_{j})\!=\!K_{H}(x_{k},y_{j})$. Putting this together with Eq. (3) into Eq. (4) yields:

To use Eq. (5) as a loss function we need to choose a kernel $K_{H}$. The most common choice of a standard multivariate normal kernel requires setting the bandwidth matrix $H$, which is non-trivial. Multivariate KDE is known to be sensitive to the optimal selection of the bandwidth matrix and the existing solutions for tuning it require solving an optimization problem which is not possible as part of a training process. Hence, we next propose an approximation, which is both practical and tractable.

3 Approximating the KL-divergence with the Contextual loss

Our approximation is based on one further observation. It is insufficient to produce an image with the same distribution of features, but rather we want additionally that the samples will be similar. That is, we ask each point $x_{i}\in X$ to be close to a specific $y_{j}\in Y$.

To achieve this, while assuming that the number of samples is large, we set the MKDE kernel such that it approximates a delta function. When the kernel is a delta, the first log term of Eq. (5) becomes a constant since:

The kernel in the second log term of Eq. (5) becomes:

We can thus simplify the objective of Eq. (5):

where we denote ${A}_{ij}=K_{H}(x_{k},y_{j})$. Next, we suggest two alternatives for the kernel ${A}_{ij}$, and show that one implies that the objective of Eq. (8) is equivalent to the Chamfer Distance , while the other implies it is equivalent to the Contextual loss of .

where $h\!>\!0$ is a scalar bandwidth parameter that we fix to $h\!=\!0.1$, as proposed in . When using these affinities our objective equals the Contextual loss of and we denote $E(x,y)=\mathcal{L}_{\text{CX}}(x,y)$.

3.2 The Chamfer Distance

A simpler way to set ${A}_{ij}$ is to take a Gaussian kernel with a fixed $H=I$ s.t. $A_{ij}=exp(-dist(x_{i},y_{j}))$. This choice implies that minimizing Eq. (8) is equivalent to minimizing the asymmetric Chamfer Distance between $X,Y$ defined as

3.3 The Contextual Loss vs. Chamfer Distance

To provide intuition on the difference between the two choices of affinity functions we present in Figure 3 an illustration in 2D. Recall, that both the Contextual loss $\mathcal{L}_{\text{CX}}$ and the Chamfer distance CD find for each point in $Y$ a single match in $X$, however, these matches are computed differently. CD selects the closet point, hence, we could get that multiple points in $Y$ are matched to the same few points in $X$. Differently, $\mathcal{L}_{\text{CX}}$ computes normalized affinities, that consider the distances between every point $x_{i}$ to all the points $y_{j}\!\in\!Y$. Therefore, it results in more diverse matches between the two sets of points, and provides a better measure of similarity between the two distributions. Additional example is presented in the supplementary.

Training with $\mathcal{L}_{\text{CX}}$ guides the network to match between the two point sets, and as a result the underlying distributions become closer. In contrast, training with CD, does not allow the two sets to get close. Indeed, we found empirically that training with CD does not converge, hence, we excluded it from our empirical reports.

Empirical analysis

The Contextual loss has been proposed in for measuring similarity between non-aligned images. It has therefore been used for applications such as style transfer, where the generated image and the target style image are very different. In the current study we assert that the Contextual loss can be viewed as a statistical loss between the distributions of features. We further assert that using such a loss during training would lead to generating images with realistic characteristics. This would make it a good candidate also for tasks where the training image pairs are aligned.

To support these claims we present in this section two experiments, and in the next section two real applications. The first experiment shows that minimizing the Contextual loss during training indeed implies also minimization of the KL-divergence. The second experiment evaluates the relation between the Contextual loss and human perception of image quality.

Since we are proposing to use the Contextual loss as an approximation to the KL-divergence, we next show empirically, that minimizing it during training also minimizes the KL-divergence.

To do this we chose a simplified super-resolution setup, based on SRResNet , and trained it with the Contextual loss as the objective. The details on the setup are provided in the supplementary. We compute during the training the Contextual loss, the KL-divergence, as well as five other common dissimilarity measures. To compute the KL-divergence, EMD and $\chi^{2}$, we need to approximate the density of each image. As discussed in Section 2, it is not clear how the multivariate solution to KDE can be smoothly used here, therefore, instead, we generate a random projection of all $5\!\times\!5$ patches onto 2D and fit them using KDE with a Gaussian kernel in 2D . This was repeated for 100 random projections and the scores were averaged over all projections (examples of projections are shown in the supplementary).

Figure 4 presents the values of all seven measures during the iterations of the training. It can be seen that all of them are minimized during the iterations. The KL-divergence minimization curve is the one most similar to that of the Contextual loss, suggesting that the Contextual loss forms a reasonable approximation.

0.2 Evaluating on Human Perceptual Judgment

Our ultimate goal is to train networks that generate images with high perceptual quality. The underlying hypothesis behind our approach is that training with an objective that maintains natural statistics will lead to this goal. To asses this hypothesis we repeated the evaluation procedure proposed in , for assessing the correlation between human judgment of similarity and loss functions based on deep features.

In it was suggested to compute the similarity between two images by comparing their corresponding deep embeddings. For each image they obtained a deep embedding via a pre-trained network, normalized the activations and then computed the $L2$ distance. This was then averaged across spatial dimension and across all layers. We adopted the same procedure while replacing the $L2$ distance with the Contextual loss approximation to KL.

Our findings, as reported in Figure 5 (the complete table is provided in the supplementary material), show the benefits of our proposed approach. The Contextual loss between deep features is more closely correlated with human judgment than $L2$ or Chamfer Distance between the same features. All of these perceptual measures are preferable over low-level similarity measures such as SSIM .

Applications

In this section we present two applications: single-image super-resolution, and high-resolution surface normal estimation. We chose the first to highlight the advantage of using our approach in concert with GAN. The second was selected since its output is not an image, but rather a surface of normals. This shows the generic nature of our approach to other domains apart from image generation, where GANs are not being used.

To asses the contribution of our suggested framework for image restoration we experiment on single-image super-resolution. To place our efforts in context we start by briefly reviewing some of the recent works on super-resolution. A more comprehensive overview of the current trends can be found in .

Recent solutions based on CNNs can be categorized into two groups. The first group relies on the $L2$ or $L1$ losses , which lead to high PSNR and SSIM at the price of low perceptual quality, as was recently shown in . The second group of works aim at high perceptual quality. This is done by adopting perceptual loss functions , sometimes in combination with GAN or by adding the Gram loss , which nicely captures textures .

Our main goal is to generate natural looking images, with natural internal statistics. At the same time, we do not want the structural similarity to be overly low (the trade-off between the two is nicely analyzed in ). Therefore, we propose an objective that considers both, with higher importance to perceptual quality. Specifically, we integrate three loss terms: (i) The Contextual loss – to make sure that the internal statistics of the generated image are similar to those of the ground-truth high-resolution image. (ii) The $L2$ loss, computed at low resolution – to drive the generated image to share the spatial structure of the target image. (iii) Finally, following we add an adversarial term, which helps in pushing the generated image to look “real”.

Given a low-resolution image $s$ and a target high-resolution image $y$, our objective for training the network $G$ is:

where in all our experiments $\lambda_{\text{CX}}\!=\!0.1$, $\lambda_{\text{GAN}}\!=\!1e-3$, and $\lambda_{{L2}}\!=\!10$. The images $G(s)^{\text{LF}},y^{\text{LF}}$ are low-frequencies obtained by convolution with a Gaussian kernel of width $21\!\times\!21$ and $\sigma\!=\!3$. For the Contextual loss feature extraction we used layer $conv3\_4$ of VGG19 .

We adopt the SRGAN architecture We used the implementation in https://github.com/tensorlayer/SRGAN and replace only the objective. We train it on just 800 images from the DIV2K dataset , for 1500 epochs. Our network is initialized by first training using only the $L2$ loss for 100 epochs.

Empirical evaluation was performed on the BSD100 dataset . As suggested in we compute both structural similarity (SSIM ) to the ground-truth and perceptual quality (NRQM ). The “ideal” algorithm will have both scores high.

Table 1 compares our method with three recent solutions whose goal is high perceptual quality. It can be seen that our approach outperforms the state-of-the-art on both evaluation measures. This is especially satisfactory as we needed only $800$ images for training, while previous methods had to train on tens or even hundreds of thousands of images. Note, that the values of the perceptual measure NRQM are not normalized and small changes are actually quite significant. For example the gap between and is only 0.014 yet visual inspection shows a significant difference. The gap between our results and is 0.08, i.e., almost 6 times bigger, and visually it is significant.

Figure 6 further presents a few qualitative results, that highlight the gap between our approach and previous ones. Both SRGAN and EnhanceNet rely on adversarial training (GAN) in order to achieve photo-realistic results. This tends to over-generate high-frequency details, which make the image look sharp, however, often these high-frequency components do not match those of the target image. The Contextual loss, when used in concert with GAN, reduces these artifacts, and results in natural looking image patches.

An interesting observation is that we achieve high perceptual quality while using for the Contextual loss features from a mid-layer of VGG19, namely $conv3\_4$. This is in contrast to the reports in that had to use for SRGAN the high-layer $conv5\_4$ for the perceptual loss (and failed when using low-layer such as $conv2\_2$). Similarly, EnhanceNet required a mixture of $pool2$ and $pool5$.

2 High-resolution Surface Normal Estimation

The framework we propose is by no means limited to networks that generate images. It is a generic approach that could be useful for training networks for other tasks, where the natural statistics of the target should be exhibited in the network’s output. To support this, we present a solution to the problem of surface normal estimation – an essential problem in computer vision, widely used for scene understanding and new-view synthesis.

The task we pose is to estimate the underlying normal map from a single monocular color image. Although this problem is ill-posed, recent CNN based approaches achieve satisfactory results on the NYU-v2 dataset . Thanks to its size, quality and variety, NYU-v2 is intensively used as a test-bed for predicting depth, normals, segmentation etc. However, due to the acquisition protocol, its data does not include the fine details of the scene and misses the underlying high frequency information, hence, it does not provide a detailed representation of natural surfaces.

Therefore, we built a new dataset of images of surfaces and their corresponding normal maps, where fine details play a major role in defining the surface structure. Examples are shown in Figure 7. Our dataset is based on 182 different textures and their respective normal maps that were collected from the Internetwww.poliigon.com and www.textures.com, originally offered for usages of realistic interior home design, gaming, arts, etc. For each texture we obtained a high resolution color image ($1024\!\times\!1024$) and a corresponding normal-map of a surface such that its underlying plane normal points towards the camera (see Figure 7). Such image-normals pairs lack the effect of lighting, which plays an essential role in normal estimation. Hence, we used Blenderwww.blender.org, a 3D renderer, to simulate each texture under $10$ different point-light locations. This resulted in a total of $1820$ pairs of image-normals. The textures were split into $90\%$ for training and $10\%$ for testing, such that the test set includes all the $10$ rendered pairs of each included texture.

The collection offers a variety of materials including wood, stone, fabric, steel, sand, etc., with multiple patterns, colors and roughness levels, that capture the appearance of real-life surfaces. While some textures are synthetic, they look realistic and exhibit imperfections of real surfaces, which are translated into the fine-details of both the color image as well as its normal map.

We propose using an objective based on a combination of three loss terms: (i) The Contextual loss – to make sure that the internal statistics of the generated normal map match those of the target normal map. (ii) The $L2$ loss, computed at low resolution, and (iii) The $L1$ loss. Both drive the generated normal map to share the spatial layout of the target map. Our overall objective is:

where, $\lambda_{\text{CX}}\!=\!1$, and $\lambda_{{L2}}\!=\!0.1$. The normal-maps $G(s)^{\text{LF}},y^{\text{LF}}$ are low-frequencies obtained by convolution with a Gaussian kernel of width $21\!\times\!21$ and $\sigma=3$. We tested with both $\lambda_{L1}\!=\!1$ and $\lambda_{L1}\!=\!0$, which removes the third term.

We chose as architecture the Cascaded Refinement Network (CRN) originally suggested for label-to-image and was shown to yield great results in a variety of other tasks . For the contextual loss we took as features $5\!\times\!5$ patches of the normal map (extracted with stride 2) and layers $conv1\_2,conv2\_2$ of VGG19. In our implementation we reduced memory consumption by random sampling of all three layers into $65\!\times\!65$ features.

Table 2 compares our results with previous solutions. We compare to the recently proposed PixelNet , that presented state-of-the-art results on NYU-v2. Since PixelNet was trained on NYU-v2, which lacks fine details, we also tried fine-tuning it on our dataset. In addition, we present results obtained with CRN . While the original CRN was trained with both $L1$ and the perceptual loss , this combination provided poor results on normal estimation. Hence, we excluded the perceptual loss and report results with only $L1$ as the loss, or when using our objective of Eq. (12). It can be seen that CRN trained with our objective (with $L1$) leads to the best quantitative scores.

We would like to draw your attention to the inferior scores obtained when removing the $L1$ term from our objective (i.e., setting $\lambda_{L1}\!=\!0$ in Eq. (12)). Interestingly, this contradicts what one sees when examining the results visually. Looking at the reconstructed surfaces reveals that they look more natural and more similar to the ground-truth without $L1$. A few examples illustrating this are provided in Figures 7, 8. This phenomena is actually not surprising at all. In fact, it is aligned with the recent trends in super-resolution, discussed in Section 4.1, where perceptual evaluation measures are becoming a common assessment tool. Unfortunately, such perceptual measures are not the common evaluation criteria for assessing reconstruction of surface normals.

Finally, we would like to emphasize, that our approach generalizes well to other textures outside our dataset. This can be seen from the results in Figure 1 where two texture images, found online, were fed to our trained network. The reconstructed surface normals are highly detailed and look natural.

Conclusions

In this paper we proposed using loss functions based on a statistical comparison between the output and target for training generator networks. It was shown via multiple experiments that such an approach can produce high-quality and state-of-the-art results. While we suggest adopting the Contextual loss to measure the difference between distributions, other loss functions of a similar nature could (and should, may we add) be explored. We plan to delve into this in our future work.

Acknowledgements the Israel Science Foundation under Grant 1089/16 and by the Ollendorf foundation.

References