Probabilistic Noise2Void: Unsupervised Content-Aware Denoising

TL;DR

Probabilistic Noise2Void (PN2V) is a self-supervised CNN-based denoising method that predicts per-pixel intensity distributions, enabling effective unsupervised image denoising across various noise conditions.

Contribution

PN2V introduces a probabilistic framework for self-supervised denoising by training CNNs to predict pixel-wise intensity distributions, improving upon existing methods.

Findings

Achieves competitive denoising results on microscopy datasets.

Works effectively across diverse noise regimes.

Outperforms traditional self-supervised methods in certain scenarios.

Abstract

Today, Convolutional Neural Networks (CNNs) are the leading method for image denoising. They are traditionally trained on pairs of images, which are often hard to obtain for practical applications. This motivates self-supervised training methods such as Noise2Void~(N2V) that operate on single noisy images. Self-supervised methods are, unfortunately, not competitive with models trained on image pairs. Here, we present 'Probabilistic Noise2Void' (PN2V), a method to train CNNs to predict per-pixel intensity distributions. Combining these with a suitable description of the noise, we obtain a complete probabilistic model for the noisy observations and true signal in every pixel. We evaluate PN2V on publicly available microscopy datasets, under a broad range of noise regimes, and achieve competitive results with respect to supervised state-of-the-art methods.

Tables

Table 1. Table 1: Results of PN2V and baseline methods on three datasets from  [ 13 ] . Comparisons are performed on five noise regimes (NR1-NR5). Numbers report PSNR (dB) ± plus-or-minus \pm 2 SEM, averaged over all 50 50 50 images in each NR. We group all supervised/non-supervised methods and mark the highest values in bold. Rows marked by asterisk ( ∗ * ) use a scale- and shift-invariant PSNR calculation to address inconsistent acquisitions in the Two-Photon mice dataset (see main text). Comp. times: All CNN based methods required below 1s per image ( NVIDIA TITAN Xp ); VST+BM3D required on avg. 6.22s.

Confocal Mice
	NR1	NR2	NR3	NR4	NR5	Mean
Input	29.38$\pm$0.01	32.44$\pm$0.01	35.59$\pm$0.01	38.90$\pm$0.01	42.64$\pm$0.03	35.79
U-Net (PN2V)	38.24$\pm$0.02	39.72$\pm$0.03	41.34$\pm$0.03	43.02$\pm$0.04	45.11$\pm$0.05	41.49
U-Net (N2V)	37.56$\pm$0.02	38.78$\pm$0.02	39.94$\pm$0.02	41.01$\pm$0.02	41.95$\pm$0.02	39.85
VST+BM3D	37.95	39.47	41.09	42.73	44.52	41.15
U-Net (trad.)	38.38$\pm$0.02	39.90$\pm$0.03	41.37$\pm$0.03	43.06$\pm$0.04	45.16$\pm$0.05	41.58
DnCNN	38.15	39.78	41.41	43.11	45.20	41.53
N2N	38.19	39.77	41.28	42.83	44.56	41.33
Confocal Zebrafish
Input	22.81$\pm$0.02	25.89$\pm$0.02	29.05$\pm$0.03	32.39$\pm$0.03	36.21$\pm$0.04	29.27
U-Net (PN2V)	32.45$\pm$0.02	33.96$\pm$0.03	35.48$\pm$0.05	37.07$\pm$0.06	39.08$\pm$0.07	35.61
U-Net (N2V)	32.10$\pm$0.02	33.34$\pm$0.03	34.43$\pm$0.04	35.39$\pm$0.04	36.21$\pm$0.03	34.30
VST+BM3D	32.00	33.75	35.30	36.78	38.32	35.23
U-Net (trad.)	32.93$\pm$0.03	34.35$\pm$0.04	35.67$\pm$0.05	37.11+0.06	39.09$\pm$0.07	35.83
DnCNN	32.44	34.16	35.75	37.28	39.07	35.74
N2N	32.93	34.37	35.71	37.06	38.65	35.74
Two-Photon Mice
Input	24.94$\pm$0.07	27.83$\pm$0.1	30.69$\pm$0.15	33.67$\pm$0.19	37.72$\pm$0.14	30.97
U-Net (PN2V)	33.67$\pm$0.33	34.58$\pm$0.39	35.42$\pm$0.42	36.58$\pm$0.37	39.78$\pm$0.24	36.00
U-Net (N2V)	33.42$\pm$0.31	34.31$\pm$0.36	35.09$\pm$0.38	36.08$\pm$0.33	37.80$\pm$0.14	35.34
VST+BM3D	33.81	34.78	35.77	36.97	39.39	36.14
U-Net (trad.)	34.35$\pm$0.19	35.32+0.23	36.14$\pm$0.27	37.48$\pm$0.27	40.28$\pm$0.2	36.72
DnCNN	33.67	34.95	36.10	37.43	40.30	36.49
N2N	34.33	35.32	36.25	37.46	39.89	36.65
$\ast$ U-Net (PN2V)	34.84$\pm$0.06	36.02$\pm$0.07	37.08$\pm$0.08	38.28$\pm$0.09	40.89$\pm$0.07	37.42
$\ast$ U-Net (N2V)	34.60$\pm$0.09	35.77$\pm$0.1	36.71$\pm$0.1	37.64$\pm$0.09	38.49$\pm$0.05	36.64
$\ast$ U-Net (trad.)	35.05$\pm$0.05	36.22$\pm$0.06	37.28$\pm$0.07	38.78$\pm$0.1	41.34$\pm$0.07	37.73

Equations

$$p(\mathbf{x}|\mathbf{s})=\prod _{i=1}^{n}p({\mathbf{x}}_i|{\mathbf{s}}_i).$$

$$\underset{\mathbf{\theta }}{\mathrm{argmin}}\sum _{i=1}^n\sum _{j=1}^m({\hat{\mathbf{s}}}_i^j-{\mathbf{s}}_i^j{)}^2$$

$$p({\mathbf{s}}_i|{\tilde{\mathbf{x}}}_{\mathrm{RF}(i)};\mathbf{\theta }).$$

$$p({\mathbf{x}}_i,{\mathbf{s}}_i|{\tilde{\mathbf{x}}}_{\mathrm{RF}(i)};\mathbf{\theta })=p({\mathbf{s}}_i|{\tilde{\mathbf{x}}}_{\mathrm{RF}(i)};\mathbf{\theta })p({\mathbf{x}}_i|{\mathbf{s}}_i).$$

$$\begin{array}{rl}p({\mathbf{x}}_i|{\tilde{\mathbf{x}}}_{\mathrm{RF}(i)};\mathbf{\theta })& ={\int }_{-\infty }^{\infty }p({\mathbf{s}}_i|{\tilde{\mathbf{x}}}_{\mathrm{RF}(i)};\mathbf{\theta })p({\mathbf{x}}_i|{\mathbf{s}}_i)d{\mathbf{s}}_i,\end{array}$$

$$\begin{array}{r}\underset{\mathbf{\theta }}{\mathrm{argmin}}\sum _{i=1}^n-\ln \left({\int }_{-\infty }^{\infty }p({\mathbf{s}}_i|{\tilde{\mathbf{x}}}_{\mathrm{RF}(i)};\mathbf{\theta })p({\mathbf{x}}_i|{\mathbf{s}}_i)d{\mathbf{s}}_i\right).\end{array}$$

$$\begin{array}{r}\underset{\mathbf{\theta }}{\mathrm{argmin}}\sum _{i=1}^n-\ln \left(\frac{1}{K}\sum _{k=1}^{K}p({\mathbf{x}}_i|{\mathbf{s}}_i^k)\right).\end{array}$$

$$\begin{array}{rl}\frac{\partial }{\partial {\mathbf{\theta }}_l}\sum _{i=1}^n-\ln \left(\frac{1}{K}\sum _{k=1}^{K}p({\mathbf{x}}_i|{\mathbf{s}}_i^k)\right)& =\sum _{i=1}^n-\left(\frac{{\sum }_{k=1}^K\frac{\partial }{\partial {\mathbf{s}}_i^k}p({\mathbf{x}}_i|{\mathbf{s}}_i^k)\frac{\partial {\mathbf{s}}_i^k}{\partial {\mathbf{\theta }}_l}}{{\sum }_{k=1}^{K}p({\mathbf{x}}_i|{\mathbf{s}}_i^k)}\right).\end{array}$$

$$\begin{array}{rl}{\mathbf{s}}_i^{\mathrm{MMSE}}& =\underset{{\hat{\mathbf{s}}}_i}{\mathrm{argmin}}{\mathbb{E}}_{p({\mathbf{s}}_i|{\mathbf{x}}_{\mathrm{RF}(i)})}\left[{\left({\hat{\mathbf{s}}}_i-{\mathbf{s}}_i\right)}^2\right]={\mathbb{E}}_{p({\mathbf{s}}_i|{\mathbf{x}}_{\mathrm{RF}(i)})}\left[{\mathbf{s}}_i\right],\end{array}$$

$$\begin{array}{rl}{\mathbf{s}}_i^{\mathrm{MMSE}}& \approx \frac{{\sum }_{k=1}^{K}p({\mathbf{x}}_i|{\mathbf{s}}_i^k){\mathbf{s}}_i^k}{{\sum }_{k=1}^{K}p({\mathbf{x}}_i|{\mathbf{s}}_i^k)}.\end{array}$$

Full text

1 Introduction

Image restoration is the problem of reconstructing an image from a corrupted version of itself. Recent work shows how CNNs can be used to build powerful content-aware image restoration (CARE) pipelines [11, 10, 12, 13, 6, 4, 1, 5]. However, for supervised CARE models, such as [10], pairs of clean and noisy images are required.

For many application areas, it is impractical or impossible to acquire clean ground-truth images [2]. In such cases, Noise2Noise (N2N) training [6] relaxes the problem, only requiring two noisy instances of the same data. Unfortunately, even the acquisition of two noisy realizations of the same image content is often difficult [2]. Self-supervised training methods, such as Noise2Void (N2V) [4], are a promising alternative, as they operate exclusively on single noisy images [4, 1, 5]. This is enabled by excluding/masking the center (blind-spot) of the network’s receptive fields. Self-supervised training assumes that the noise is pixel-wise independent and that the true intensity of a pixel can be predicted from local image context, excluding before-mentioned blind-spots [4]. For many applications, especially in the context of microscopy images, the first assumption is fulfilled, but the second assumption offers room for improvements [5].

Hence, self-supervised models can often not compete with supervised training [4]. In concurrent work, by Laine et al. [5], this problem was elegantly addressed by assuming a Gaussian noise model and predicting Gaussian intensity distributions per pixel. The authors also showed that the same approach an be applied to other noise distributions, which can be approximated as Gaussian, or can be described analytically.

Here, we introduce a new training approach called Probabilistic Noise2Void (PN2V). Similar to [5], PN2V proposes a way to leverage information of the network’s blind-spots. However, PN2V is not restricted to Gaussian noise models or Gaussian intensity predictions. More precisely, to compute the posterior distribution of a pixel, we combine $(i)$ a general noise model that can be represented as a histogram (observation likelihood), and $(ii)$ a distribution of possible true pixel intensities (prior), represented by a set of predicted samples.

Having this complete probabilistic model for each pixel, we are now free to chose which statistical estimator to employ. In this work we use MMSE estimates for our final predictions and show that MMSE-PN2V consistently outperformes other self-supervised methods and, in many cases, leads to results that are competitive even with supervised state-of-the-art CARE networks (see below).

2 Background

Image Formation and the Denoising Task:

An image $\boldsymbol{x}=(\boldsymbol{x}_{1},\dots,\boldsymbol{x}_{n})$ is the corrupted version of a clean image (signal) $\boldsymbol{s}=(\boldsymbol{s}_{1},\dots,\boldsymbol{s}_{n})$. Our goal is to recover the original signal from $\boldsymbol{x}$, thus implementing a function $f(\boldsymbol{x})=\hat{\boldsymbol{s}}\approx\boldsymbol{s}$.

In this paper, we assume that each observed pixel value $\boldsymbol{x}_{i}$ is independently drawn from the conditional distribution $p(\boldsymbol{x}_{i}|\boldsymbol{s}_{i})$ such that

[TABLE]

We will refer to $p(\boldsymbol{x}_{i}|\boldsymbol{s}_{i})$ as observation likelihood. It is described by an arbitrary noise model.

Traditional Training and Noise2Noise:

The function $f(\boldsymbol{x})$ can be implemented by a Fully Convolutional Network (FCN) [7] (see e.g. [11, 10, 12, 6]), a type of CNN that takes an image as input and produces an entire (in this case denoised) image as output. However, in this setup every predicted output pixel $\hat{\boldsymbol{s}}_{i}$ depends only on a limited receptive field $\boldsymbol{x}_{\mathrm{RF}(i)}$, i.e. a patch of input pixels surrounding it. FCN based image denoising in fact implements $f(\boldsymbol{x})$ by producing independent predictions $\hat{\boldsymbol{s}}_{i}=g(\boldsymbol{x}_{\mathrm{RF}(i)};\boldsymbol{\theta})\approx\boldsymbol{s}_{i}$ for each pixel $i$, depending only on $\boldsymbol{x}_{\mathrm{RF}(i)}$ instead of on the entire image. The prediction is parametrized by the weights $\boldsymbol{\theta}$ of the network.

In traditional training, $\boldsymbol{\theta}$ are learned from pairs of noisy $\boldsymbol{x}^{j}$ and corresponding clean training images $\boldsymbol{s}^{j}$, which provide training examples $(\boldsymbol{x}_{\mathrm{RF}(i)}^{j},\boldsymbol{s}_{i}^{j})$ consisting of noisy input patches $\boldsymbol{x}_{\mathrm{RF}(i)}^{j}$ and their corresponding clean target values $\boldsymbol{s}_{i}^{j}$. The parameters $\boldsymbol{\theta}$ are traditionally tuned to minimize an empirical risk function such as the average squared distance

[TABLE]

over all training images $j$ and pixels $i$.

In Noise2Noise [6], Lehtinen et al. show that clean data is in fact not necessary for training and that the same training scheme can be used with noisy data alone. Noise2Noise uses pairs of corresponding noisy training images $\boldsymbol{x}^{j}$ and ${\boldsymbol{x}^{\prime}}^{j}$, which are based on the same signal $\boldsymbol{s}^{j}$, but are corrupted independently by noise (see Eq. 1). Such pairs can for example be acquired by imaging a static sample twice. Noise2Noise uses training examples $(\boldsymbol{x}_{\mathrm{RF}(i)}^{j},{\boldsymbol{x}^{\prime}}^{j}_{i})$, with the input patch $\boldsymbol{x}_{\mathrm{RF}(i)}^{j}$ cropped from the first image $\boldsymbol{x}^{j}$ and the noisy target ${\boldsymbol{x}^{\prime}}^{j}_{i}$ extracted from the patch center in the second one ${\boldsymbol{x}^{\prime}}$. It is of course impossible for the network to predict the noisy pixel value ${\boldsymbol{x}^{\prime}}^{j}_{i}$ from the independently corrupted input $\boldsymbol{x}_{\mathrm{RF}(i)}^{j}$. However, assuming the noise is zero centered, i.e. $\mathbb{E}\left[{\boldsymbol{x}^{\prime}}_{i}^{j}\right]=\boldsymbol{s}_{i}^{j},$ the best achievable prediction is the clean signal $\boldsymbol{s}_{i}^{j}$ and the network will learn to denoise the images it is presented with.

Noise2Void Training:

In Noise2Void, Krull et al. [4] show that training is still possible when not even noisy training pairs are available. They use single images to extract input and target for their networks. If this was done naively, the network would simply learn the identity transformation, directly outputting the value at the center of each pixel’s receptive field. Krull et al. address the issue by effectively removing the central pixel from the networks receptive field. To achieve this, they mask the pixel during training, replacing it with a random value from the vicinity. Thus, a Noise2Void trained network can be seen as a function $\hat{\boldsymbol{s}}_{i}=\tilde{g}(\tilde{\boldsymbol{x}}_{\mathrm{RF}({i})};\boldsymbol{\theta})\approx\boldsymbol{s}_{i},$ making a prediction for a single pixel based on the modified patch $\tilde{\boldsymbol{x}}_{\mathrm{RF}({i})}$ that excludes the central pixel. Such a network can no longer describe the identity, and can be trained from single noisy images.

However, this ability comes at price. The accuracy of the predictions is reduced, as the network has to exclude the central pixel of its receptive field, thus having less information available.

To allow efficient training of a CNN with Noise2Void, Krull et al. simultaneously mask multiple pixels in larger training patches and jointly calculate their gradients.

3 Method

Maximum Likelihood Training:

In PN2V, we build on the idea of masking pixels [4] to obtain a prediction from the modified receptive field $\tilde{\boldsymbol{x}}_{\mathrm{RF}({i})}$. However, instead of directly predicting an estimate for each pixel value, PN2V trains a CNN to describe a probability distribution

[TABLE]

We will refer to $p(\boldsymbol{s}_{i}|\tilde{\boldsymbol{x}}_{\mathrm{RF}({i})};\boldsymbol{\theta})$ as prior, as it describes our knowledge of the pixel’s signal considering only its surroundings, but not the observation at the pixel itself $\boldsymbol{x}_{i}$, since it has been excluded from $\tilde{\boldsymbol{x}}_{\mathrm{RF}({i})}$. We choose a sample based representation for this prior, which will be discussed below.

Remembering that the observed pixels values are drawn independently (Eq. 1), we can combine Eq. 3 with our noise model, and obtain the joint distribution

[TABLE]

By integrating over all possible clean signals, we can derive

[TABLE]

the probability of observing the pixel value $\boldsymbol{x}_{i}$, given we know its surroundings $\tilde{\boldsymbol{x}}_{\mathrm{RF}({i})}$. We can now view CNN training as an unsupervised learning task. Following the maximum likelihood approach, we tune $\boldsymbol{\theta}$ to minimize

[TABLE]

Note that in order to improve readability, we from here on omit the index $j$, and refrain from explicitly referring to the training image.

Sample Based Prior:

To allow an efficient optimization of Eq. 6 we choose a sample based representation of our prior $p(\boldsymbol{s}_{i}|\tilde{\boldsymbol{x}}_{\mathrm{RF}({i})};\boldsymbol{\theta})$. For every pixel $i$, our network directly predicts $K=800$ output values $\boldsymbol{s}^{k}_{i}$, which we interpret as independent samples, drawn from $p(\boldsymbol{s}_{i}|\tilde{\boldsymbol{x}}_{\mathrm{RF}({i})};\boldsymbol{\theta})$. We can now approximate Eq. 6 as

[TABLE]

During training we use Eq. 7 as loss function. Note that the summation over $k$ can be efficiently performed on the GPU.

Since every sample $\boldsymbol{s}^{k}_{i}$ is effectively a function of the parameters $\boldsymbol{\theta}$, we can calculate the derivative with respect to any network parameter $\boldsymbol{\theta}_{l}$ as

[TABLE]

Minimal Mean Squared Error (MMSE) Inference:

Assuming our network is sufficiently trained, we are now interested in processing images and finding sensible estimates for every pixel’s signal $\boldsymbol{s}_{i}$. Based on our probabilistic model, we derive the MMSE estimate, which is defined as

[TABLE]

where $p(\boldsymbol{s}_{i}|\boldsymbol{x}_{\mathrm{RF}(i)})$ is the posterior distribution of the signal given the complete surrounding patch. The posterior is proportional to the joint distribution given in Eq. 4. We can thus approximate $\boldsymbol{s}^{\mathrm{MMSE}}_{i}$ by weighing our predicted samples with the corresponding observation likelihood and calculating their average

[TABLE]

Figure 1 illustrates the process and shows the involved distributions for a concrete pixel in a real example.

4 Experiments

The results of our experiments can be found in Table 1. In Figure 2 we provide qualitative results on realistic test images.

Datasets:

We evaluate PN2V on datasets provided by Zhang et al. in [13]. Since PN2V is not yet implemented for multi-channel images, we use all available single-channel datasets.

These datasets are recorded with different samples and under different imaging conditions. Each of them consists of a total of 20 fields of view (FOVs). One FOV is reserved for testing. The other 19 are used for training and validation.

For each FOV, the data is composed of 50 raw microscopy images, each containing different noise realizations of the same static sample. For every FOV, Zhang et al. additionally simulate four reduced noise regimes (NRs) by averaging different subsets of 2, 4, 8, and 16 raw images [13]. We will refer to the raw images as NR1 and to the regimes created through averaging 2, 4, 8, and 16 images as NR2, NR3, NR4, and NR5, respectively.

We find that in one of the datasets (Two-Photon Mice) the average pixel intensity fluctuates heavily over the course of the 50 images, even though it should be approximately constant for each FOV. Considering that a single ground truth image (the average) is used for the evaluation on all 50 images, this leads to fluctuations and distortions in the calculated PSNR values, which are also reflected in the comparatively high standard errors (SEMs) for all methods, see Table 1. To account for this inconsistency in the data, we additionally use a variant of the PSNR calculation that is invariant to arbitrary shifts and linear transformations in the ground truth signal. These values are marked by an asterisk (*). Details can also be found in the supplementary material111 The Supplement will be made available soon..

Acquiring Noise Models:

In our experiments, we use a histogram based method to measure and describe the noise distribution $p(\boldsymbol{x}_{i}|\boldsymbol{s}_{i})$. We start with corresponding pairs of clean $\boldsymbol{s}^{j}$ and noisy $\boldsymbol{x}^{j}$ images. Here, we use the available training data from [13] for this purpose. However, in general these images could show an arbitrary test pattern that covers the desired range of values and do not have to resemble the sample we are interested in. We construct a 2D histogram (256$\times$256 bins), with the y- and x-axis corresponding to the clean $\boldsymbol{s}^{j}_{i}$ and noisy pixel values $\boldsymbol{x}^{j}_{i}$, respectively. By normalizing every row, we obtain a a probability distribution for every signal. Considering Eq. 7, we require our model to be differentiable with respect to the $\boldsymbol{s}_{i}$. To ensure this differentiability, we linearly interpolate along the y-axis of the normalized histogram, obtaining a model for $p(\boldsymbol{x}_{i}|\boldsymbol{s}_{i})$ that is continuous in $\boldsymbol{s}_{i}$.

Evaluated Denoising Models/Methods:

To put the denoising results of PN2V into perspective, we compare to various state-of-the-art baselines, including the strongest published numbers on the datasets.

U-Net (PN2V):

We use a standard U-Net [9]. Our network has a U-Net depth of 3, 1 input channel, and $K=800$ output channels, which are interpreted as samples. We use a initial feature channel number of 64 in the fist U-Net layer. We train our network separately for each NR in each dataset. We use the same masking technique as [4]. Further details and training parameters can be found in the supplementary material11footnotemark: 1.

U-Net (N2V):

We use the same network architecture as for U-Net (PN2V) but modify the outputlayer to produce only a single prediction instead of $K=800$. The network is trained using the Noise2Void scheme as described in [4]. All training parameters are identical to U-Net (PN2V).

U-Net (trad.):

We use the exact same architecture as U-Net (N2V), but train the network using the available ground-truth data and the standard MSE loss (see Eq. 2). All training parameters are identical to U-Net (PN2V) and U-Net (N2V).

VST+BM3D:

Numbers are taken from [13]. The authors fit a Poisson Gaussian noise model to the data and then apply a combination of variance-stabilizing transformation (VST) [8] and BM3D filtering [3].

DnCNN:

Numbers are taken from [13]. DnCNN [12] is an established CNN based denoising architecture that is trained in a supervised fashion.

N2N:

Numbers are taken from [13]. The authors train a network according to the N2N scheme, using an architecture similar to the one presented in [6].

5 Discussion

We have introduced PN2V, a fully probabilistic approach extending self-supervised CARE training. PN2V makes use of an arbitrary noise model which can be determined by analyzing any set of available images that are subject to the same type of noise. This is a decisive advantage compared to state-of-the-art supervised methods and allows PN2V to be used for many practical applications.

The much improved performance of PN2V lies consistently beyond self-supervised training and can often compete with state-of-the-art supervised methods. We see a plethora of unique applications for PN2V, for example in challenging low-light conditions, where noise typically is the limiting factor for downstream analysis.

Acknowledgements

We thank Uwe Schmidt, Martin Weigert, and Tobias Pietzsch for the helpful discussions. We thank Tobias Pietzsch for proof reading. The computations were performed on an HPC Cluster at the Center for Information Services and High Performance Computing (ZIH) at TU Dresden.