Image Restoration with Locally Selected Class-Adapted Models

Afonso M. Teodoro, José M. Bioucas-Dias, Mário A. T. Figueiredo

Introduction

Denoising is one of the oldest and most central problems in image processing, dating back to the early days of the field . Although it has been argued that the current best generic methods are very close to the theoretically maximum possible performance , image denoising is still a very active area of research. Most, if not all, state-of-the-art methods belong to the patch-based family, i.e., they process the noisy image on a patch-by-patch fashion, relying on techniques such as non-local means , collaborative filtering , dictionary learning , or statistical models of the patches , , , . In more general imaging inverse problems, such as deblurring, computed tomography, or magnetic resonance imaging, to mention only a few classical examples, it may not be obvious how these patch-based techniques may be applied, and this has recently been a topic of interest , , , .

Whereas the bulk of the work in image restoration aims at developing methods of general applicability, there has been some recent work on developing class-specific methods , . As the name suggests, these methods are tailored to perform very well on a certain class of images, for example, text, faces, or some type of medical images. Indeed, if we knew the class of the image being processed/reconstructed, it would be expectable that a targeted method could outperform a generic one, i.e., one that does not take into account the specific class of the image in hand. Whereas in some cases, it is known that the image being estimated belongs to a certain class (e.g., brain MR or CT images, face images, fingerprints, text), in many situations this is not the case and there may even be regions of different classes present in the image (for example, a document image may contain text and one or more faces or some other type of images). This second type of scenario (where several classes may be present in a given image) is the one addressed in this paper. Identifying the classes that are present at each location of an image can be interpreted as a form of image segmentation, thus the problem formulated and addressed in this paper may be seen as that of simultaneous image restoration and segmentation.

Recently , we have proposed a method for class-adapted restoration/reconstruction, which builds upon the so-called plug-and-play approach , by plugging a class-adapted denoiser based on Gaussian mixture models (GMM) into the iterations of an alternating direction method of multipliers (ADMM) algorithm. Experiments reported in (both in deblurring and compressive imaging) have shown that the proposed method yields state-of-the-art results, when applied to images known to contain text or a face, clearly outperforming the best generic techniques, such as IDD-BM3D .

In this paper, we extend the plug-and-play class-adapted approach to the scenario mentioned above: the image being restored is from an unknown class and/or may contain regions from different classes (e.g., text and faces, or text and natural images). As mentioned above, this may be seen as a method to perform simultaneous segmentation and restoration, by exploiting synergies between these two tasks. Notice that our main goal is not to obtain a good or meaningful segmentation, but to use the segmentation to allow class-specific models to be exploited at each location of the image, thus our focus is on the restoration performance of the method. Nevertheless, we believe that this type of approach may contribute to bridging the gap between low-level (restoration) and mid-level (segmentation) image processing/analysis.

The remaining sections of the paper are organized as follows. Section 2 reviews the classical image restoration formulation and the basic tools upon which our method is built: ADMM, the plug-and-play scheme, and patch-based image denoising using GMM priors. The proposed method is described in Section 3, while experimental results are reported in Section 4. Finally, Section 5 concludes the paper and gives pointers for future work.

Formulation and Tools

The classical formulation of image reconstruction/restoration problems is

2 ADMM and Plug-and-Play

In recent years, variable splitting algorithms, such as ADMM, have received a lot of attention, in particular in the image processing and machine learning communities . A notable feature of ADMM, as applied to (2), is that it separates the handling of the data term (log-likelihood) from that of the prior/regularizer. The standard instantiation of ADMM to tackle (2) consists in the cyclic application of the following steps , :

Problem (3) is quadratic and has a closed-form solution, which requires solving a linear system (equivalently, inverting a matrix):

Although this is sometimes seen as an obstacle (motivating, for example, the introduction of linearized versions of ADMM ), there are cases in which this inversion can be computed very efficiently using fast transforms (namely the FFT) , , and in those cases ADMM exhibits state-of-the-art speed. Image deconvolution (with periodic or other boundary conditions ) and inpainting are two of those cases where ADMM excels.

and can be seen as the MAP solution of a denoising problem, where the argument of $\mbox{prox}_{\phi}$ is the noisy data, the noise is Gaussian i.i.d. with unit variance, and the prior is $p({\bf x})\propto\exp(-\phi(\bf x))$ .

In recent work, it has been suggested that the denoiser corresponding to the MPO in (4) can be replaced with an off-the-shelf state-of-the-art denoising algorithm, such as BM3D ; that approach has been termed plug-and-play . Since, in general, a denoising algorithm may not correspond necessarily to the MPO of a convex regularizer/log-prior, convergence of the resulting algorithm requires further analysis; initial results have been presented in .

3 Patch-Based Denoising with GMM Priors

It has been shown that a simple GMM, estimated from a collection of clean images, is a remarkably effective patch prior . More recently, we have shown that excellent denoising performance is also obtained if this GMM is estimated directly from the patches of the noisy image to be denoised, and that the corresponding expectation-maximization (EM) algorithm is a simple modification of the one for estimating the GMM from noiseless patches . Moreover, since the prior is a GMM, it is easy to compute the conditional expectation of each patch given its noisy version, which is the optimal MMSE estimate (in contrast with , where a MAP estimate is used). Letting ${\bf x}_{i}$ and ${\bf y}_{i}$ denote an arbitrary patch of images $\bf x$ and $\bf y$ (unknown clean image and noisy one, respectively), the probabilistic model for patch denoising is

where $\mathcal{N}(\cdot;\mbox{\boldmath$ \mu $},{\bf C})$ denotes a Gaussian probability density function of mean $\mu$ and covariance $\bf C$ . The resulting MMSE estimate of ${\bf x}_{i}$ is given by

Notice that $\beta_{m}({\bf y}_{i})$ is simply the posterior probability that the $i$ -th patch was generated by the $m$ -th GMM component, whereas ${\bf v}_{m}({\bf y}_{i})$ is the conditional MMSE (and MAP) estimate of ${\bf x}_{i}$ , if we knew that it had been generated by the $m$ -th GMM component.

The final denoised image is assembled by putting the patch estimates back in their locations. Since the patches overlap, there are several estimates of each pixel, which are usually combined by straight averaging. In , we proposed to use the optimal weighted averaging, where the weights are the inverses of the posterior variances, which can also be computed in closed-form. The use of weights based on the posterior variance of the patch estimates was also used in , but with a single Gaussian per patch.

Recently , we have proposed to use GMM patch-based denoising in a plug-and-play approach, to deblur images of specific classes. For that purpose, the GMM is estimated from a collection of clean images of the class of interest, and the denoising step (4) of the ADMM algorithm is replaced with the GMM-patch-based method described above. Experiments reported in show that the method produces excellent results in deconvolving text and face images, outperforming the generic state-of-the-art method IDD-BM3D .

Proposed Method

In this paper, we extend the method proposed in to handle images of unknown classes or even containing regions from different classes. Rather than a single class-adapted GMM, estimated from a collection of clean images from that class, consider $C$ classes, each of which modelled by a GMM,

where $c_{i}\in\{1,...,C\}$ is the class label of the $i$ -the patch, and $C$ the total number of classes. To estimate the $i$ -th patch, we begin by classifying it into one of the classes, and then use the corresponding GMM to obtain an MMSE estimate of that patch (given by (9)), conditioned on its noisy version. We emphasize that this process is repeated until some stopping criterion is met.

As mentioned in Section 1, classifying each patch into one of the classes can be seen as performing image segmentation. However, we stress again that segmentation is not the main goal of the proposed approach, thus we will only focus on its performance in terms of denoising and deblurring. To classify the patches, we consider two alternatives:

Simply classifying each patch independently using the maximum-likelihood criterion,

Jointly classifying all the patches under a Markov random field prior $p({\bf c})$ (where ${\bf c}$ denotes the field of all the patch class labels), more specifically a Potts prior ,

To solve (14) we use the $\alpha$ -expansion graph-cut algorithm proposed in . For more details about MRF priors for image segmentation, see , .

In summary, the multi-class denoiser described in the two preceding paragraphs is used in the plug-and-play approach described in Subsection 2.2 (see Algorithm 1).

Experimental Results

We start by presenting some results on image denoising. Table 1 compares the results obtained with the denoising algorithm, with and without classification of the patches. As a baseline, we also present the results of a state-of-the-art denoising algorithm, BM3D (with default parameters). In every run, the patch size was set to 8 by 8, and we trained a GMM with 20 components on the noisy patches, using the approach described in . Furthermore, the following classes were considered: text, faces, brain MRI, fingerprints, and generic. All of the external GMMs, also with 20 components each, were trained with samples from the corresponding class, except the generic GMM, which was trained using random images from the Berkeley dataset for image segmentation (BSDS300) .

We conclude that the proposed modification does not have a significant impact on the denoising performance. Arguably, this is due to the fact that, in pure denoising, it is possible to estimate a GMM from the noisy image itself . The resulting model is thus better adapted to the input image than if it would be trained from a different set of images, even if these images are from the same class.

Figure 1 illustrates the difference in the patch labelling, if we consider classifying each patch independently via the maximum-likelihood (ML) criterion or using the $\alpha$ -expansion graph-cut algorithm. While we do not expect the latter labelling to perform better in terms of denoising or deblurring, it is more meaningful from a segmentation viewpoint, since it is exhibits higher spatial coherence. In this example, we used three different models: one trained from the noisy image itself (black), one targeted to text images (grey), and one trained with generic images (white).

In our deblurring experiments, we considered the same classes, where each GMM (with 20 components) was trained using $6\times 6$ patches. Table 2 shows the results of ADMM-GMM algorithm, with and without classification, for all blur kernels in . As a benchmark, we also present the results obtained with state-of-the-art IDD-BM3D (with default parameters). In this set up, we assumed that the class of the input image is unknown a priori, even in the text or face images so, when no classification is done, the algorithm uses only the generic GMM. Otherwise, we let the algorithm decide which class should be used for each patch, which explains the discrepancy that exists relative to the results reported in . Furthermore, after 100 iterations of the algorithm, we switched the generic GMM with another GMM trained from the deblurred patches, which we assume to be reasonable estimates of the clean patches by then.

Since we cannot learn a GMM from the blurred input image, the improvement that is achieved with the proposed method is much more visible when, in fact, the input image contains one or more of the considered classes. On the one hand, regarding the Cameraman and House images (from the generic class), we observe that the proposed scheme performs worse than IDD-BM3D ( $0.76$ dB in the worst-case scenario), while using classification may or may not improve the performance. On the other hand, for the remaining examples, ADMM-GMM achieves better results than the generic IDD-BM3D. Although the parameters of IDD-BM3D were not tuned for these examples, we emphasize that when comparing ADMM-GMM with and without classification, the classification scheme that was proposed in this paper consistently improves the results on the images that contain one or more of the considered classes.

Figure 2 shows the patch labelling at different iterations. At first, as one would expect, the classes are not correctly identified; as the algorithm progresses, the labelling becomes more accurate. In this example we used six different classes but, for visualization purposes, we display only three: face (white), text (grey), and other (black).

Conclusions and Future work

Recent work , has shown that class-adapted image priors are able to outperform generic ones, when the input image in fact contains one or more of the considered classes. In this paper, we developed a method that automatically identifies which patches of the observed image should be reconstructed using class-adapted priors. This is an important feature for two main reasons: first, we often do not know whether the input image is from a particular class or not; second, we may have more than one class in a single image.

Several aspects of the proposed approach still need improvement. First, although there is some very recent work regarding the convergence of ADMM with plug-and-play denoisers, there is a need to carefully analyse the convergence of ADMM with the GMM-based denoiser. Indeed, convergence of the algorithm is observed in practice, but theoretical support is still not available. Second, in the experiments reported in Section 4, parameter $\mu$ was hand-tuned; in future work, we will pursue more sophisticated techniques, such as , to adjust this parameter in an automatic way.

Finally, in this paper we tested only a few very distinct image classes, but using more pre-computed models could eventually lead to better results. Some examples of possible classes are textures, sky, buildings, and so on.