Zero-Shot Image Denoising for High-Resolution Electron Microscopy

In recent years, supervised image denoising methods based on deep neural networks e.g., DnCNN [8] have achieved great success and outperform conventional image denoiser. However, the acquisition of aligned noisy-clean images is infeasible and impractical in many scientific imaging applications, such as electron microscopy, which limits the use of supervised deep learning approaches.

Noise2Noise (N2N) [25] is the first work that proposes training a deep denoiser using paired noisy images of the same scene and demonstrates that it is statistically equivalent to supervised learning. Subsequently, self-supervised denoising methods have been proposed, enabling the denoiser to be trained from individual noisy images without paired noisy images. Noise2Void (N2V) [12] and Noise2Self (N2S) [11] design the masked-based blind-spot network (BSN) for self-supervised learning. Specifically, the masked-based BSN replaces certain pixels of input noisy images and predicts their value based on neighboring pixels. However, masked-based BSN suffers from the limited pixels of supervision and imperfect replacement strategies, leading to inefficient training and artifacts in the denoised results [13, 16]. To address these issues, the following works [13, 20, 14, 17, 26] have proposed to design novel BSN architectures by incorporating centrally masked convolution and dilated convolutions. However, these methods are limited by the large amount of calculation and the inflexible network structure. Different from BSN methods, another category of self-supervised learning approaches proposed generating approximate paired noisy images from individual noisy images. For example, Noisier2Noise [18] generates paired noisy images by adding additional noise to noisy images, and similar ideas are explored in [27, 23]. However, these methods require a known noise model of original noisy images to guarantee denoising performance. Additionally, Neighbor2Neighbor (NB2NB) [19] proposed to sub-sample the neighbor pixels of the noisy image to generate paired noisy images. Since a gap exists between the underlying ground truth of paired neighbor sub-sampled images, NB2NB has limited performance in low-SNR scenarios.

Deep image prior (DIP) [28] first demonstrated that convolutional neural networks (CNNs) inherently learn image priors with their inductive bias. DIP proposed using early-stopping strategies for recovering corrupted images with CNNs. Subsequently, several works leverage various approaches such as denoising [29], GANs [30] or diffusion models [31] to learn the prior image distribution for zero-shot image restoration. However, these approaches require a substantial amount of data for pre-training to learn the data distribution, which is impractical in the context of HREM. Therefore, we focus on zero-shot learning restoration methods, training with only a single corrupted image without additional external data. Compared with dataset-based training, zero-shot learning is significantly more challenging to avoid overfitting. In image dehazing, several works [32, 33] have achieved remarkable performance by training on a single hazy image. For zero-shot image denoising, Self2Self [15] incorporates a dropout regularization with a blind-spot networks framework to mitigate overfitting. FBI-Denoiser [17] showed superior performance by carefully designing the BSN. ZS-Noise2Noise [22] introduced a lightweight network to achieve comparable performance in zero-shot denoising. Chen et al. [34] proposed a zero-shot method for medical image artifact reduction.

Conventional spatial filter based methods have been applied to EM, such as Bilateral filter, Non-local Means, BM3D, etc. Recent, deep learning based methods have been applied in EM imaging[35, 36]. Refs. [37, 38] proposed using Cycle-GAN for STEM images denoising without paired training images. Chong et al. [39] utilized paired noisy images achieving real-world optical and electron microscopy data denoising. Refs. [40] [41] proposed the simulation-based denoising (SBD) framework creating large simulated datasets to train CNNs for denoising in a supervised learning manner. Mohan et al. [41] introduced simulated paired clean/noisy HREM images of the same substance for various imaging parameters. However, creating large simulated datasets can be a time-consuming process, and its performance will degrade sharply due to the domain gap between simulated and real noisy images. Thus, it is imperative to propose a zero-shot image denoising method that is robust to ultra-low SNR to enhance HREM data quality effectively.

We propose training a denoiser using pairs of noisy images with different resolutions, generated from an individual image. These paired images consist of a sub-sampled noisy image $S(\mathbf{y})$ and its corresponding full-resolution noisy image $\mathbf{y}$. Following the N2N strategy, the self-supervised loss can be generally stated as

However, directly minimizing the loss function above may result in the network learning an identity mapping for the noisy input pixels. Thus, we proposed to use unsampled noisy pixels for network supervision. Under this supervision, we can theoretically prove that using noisy images for network training is equivalent to using clean signals for supervision.

Based on the $\mathcal{J}$-invariant theory, the sub-sampling operation partitions the image with $n$ pixels into two disjoint sets denoted as $J$ and $J^{\complement}$, where $|J|+|J^{\complement}|=n$. Consequently, the sub-sampled noisy image $S(\mathbf{y})$ can be represented as $\mathbf{y}_{J^{\complement}}$, while the unsampled noisy pixels $\mathbf{y}\setminus S(\mathbf{y})$ are denoted as $\mathbf{y}_{J}$. When training a network using paired noisy images with different resolutions, the resulting network naturally becomes a $\mathcal{J}$-invariant function. This is because the output $f_{\theta}(\mathbf{y}_{J^{\complement}})_{J}$ does not rely on the specific values of $\mathbf{y}_{J}$.

Self-supervised image denoising methods typically assume that noise is signal-independent and spatially uncorrelated. Therefore, it is crucial to sub-sample a noisy image to maintain the denoising performance when training with paired noisy images of different resolutions. Recently, pixel-shuffling downsampling (PD) [42, 26] has emerged as an effective technique for breaking the spatial correlation of real-world noise. Motivated by the success of PD, we introduce a Random Sub-sampler for generating sub-sampled noisy images. The Random Sub-sampler follows a similar sub-sampling strategy as PD but introduces additional randomness into the sub-sampling operation. The randomness serves as a data augmentation strategy, which helps to prevent training overfitting.

The process of using Random Sub-sampler with sampling stride $s$ to generate a sub-sampled noisy image $\mathbf{y}_{J^{\complement}}\in\mathbb{R}^{H/s\times W/s}$ from an image $\mathbf{y}\in\mathbb{R}^{H\times W}$ is summarized as follows:

1. 1.

   Perform an inverse pixel shuffling (PS) [43] operation on image $\mathbf{y}\in\mathbb{R}^{H\times W}$ with a stride of $s$, resulting in: $\text{PS}_{s}^{-1}(\mathbf{y})\in\mathbb{R}^{H/s\times W/s\times s^{2}}$;

2. 2.

   For $(i,j)$-th location of sub-sampled image $\mathbf{y}_{J^{\complement}}$, randomly select one elements from ${\text{PS}_{s}^{-1}(\mathbf{y})}_{ij}$;

3. 3.

   Generate a binary matrix mask $\mathbf{m}_{J}\in\mathbb{R}^{H\times W}$ to select unsampled pixels in the original image $\mathbf{y}$. The mask $\mathbf{m}_{J}$ is set to 0 at locations where an element from $\mathbf{y}$ is selected, and 1 otherwise.

Following this process, the sub-sampled noisy image $\mathbf{y}_{J^{\complement}}$ with dimensions of $H/s\times W/s$ is obtained. The binary matrix mask $\mathbf{m}_{J}$ is also generated to identify the unsampled pixels in $\mathbf{y}$. Fig. 2 illustrates the workflow of the Random Sub-sampler, demonstrating the generation of a sub-sampled image from an input image of size $4\times 4$ with a stride of $s=2$.

The stride $s$ determines the relatively sub-sampling interval and sampling ratio of noisy pixels, which have an impact on the training of Noise2SR. In Section IV-F, we conduct a comprehensive study to evaluate the influence of the sampling stride $s$ of the Random Sub-sampler. We also compare the performance between fix-location and random-location sub-sampling strategies.

Given a noisy image $\mathbf{y}\in\mathbb{R}^{H\times W}$, We can parameterize the Noise2SR as CNN-based denoising and SR function $f_{\theta}:\mathbb{R}^{|J^{\complement}|}\to\mathbb{R}^{H\times W}$, where the parameters $\theta$ are weights of the network, $|J^{\complement}|$ is number of sub-sampled noisy image pixels. Noise2SR takes the sub-sampled noisy image $\mathbf{y}_{J^{\complement}}$ as input and outputs a prediction of denoised image $\hat{\mathbf{x}}=f_{\theta}(\mathbf{y}_{J^{\complement}})\in\mathbb{R}^{H\times W}$.

Based on the proof provided in Section III-B, we optimize the network by minimizing the loss function that compares the unsampled pixels of the original noisy image denotes as $\mathbf{y}_{J}$, with corresponding pixels in the network output, denotes as $f_{\theta}(\mathbf{y}_{J^{\complement}})_{J}$. To facilitate the selection of these pixels, we can utilize a binary mask $\mathbf{m}_{J}$. Specifically, the unsampled pixels of the original image can be obtained by element-wise multiplication with the mask, i.e., $\mathbf{y}_{J}=\mathbf{y}\odot\mathbf{m}_{J}$. Similarly, the corresponding pixels in the network output can be obtained as $f_{\theta}(\mathbf{y}_{J^{\complement}})_{J}=f_{\theta}(\mathbf{y}_{J^{\complement}})\odot\mathbf{m}_{J}$. Here, the symbol $\odot$ represents the Hadamard product. Thus, the Noise2SR network can learn a denoising and SR function by minimizing the self-supervised loss

where $\mathcal{J}$ represents the set of partitions used during the training stage, and $N$ denotes the number of images in the dataset.

The pipeline of clean image reconstruction using the well-trained N2SR model is shown in Fig. 1.B. For a random sub-sampled noisy image $\mathbf{y}_{J^{\complement}}$, well-trained Noise2SR can generated a plausible denoised result $\hat{\mathbf{x}}=f_{\hat{\theta}}(\mathbf{y}_{J^{\complement}})$. Since the sub-sampled image is generated from the original noisy image randomly, the clean image can be expressed as:

where $\mathcal{J}$ is the possible Random Sub-sampler partition set. Since the randomness of sub-sampling operation , $|\mathcal{J}|$ is numerous and $p(\mathbf{x}|\mathbf{y}_{J^{\complement}})$ is undetermined, making it challenging to precisely compute $\mathbb{E}(\hat{\mathbf{x}})$.

Thus, we propose to use MMSE estimation to approximate $\mathbb{E}(\hat{{\mathbf{x}}})$ of the clean signal. Given a set of sub-sampled noisy images $\mathbf{y}_{J^{\complement}}$, we can approximate the MMSE estimate of the clean image by averaging all the plausible denoised results $f_{\hat{\theta}}(\mathbf{y}_{J^{\complement}})$. Specifically, we first randomly sub-sampled the noisy image $M$ times to obtain a sub-sampled noisy image set $\mathcal{Y}=\left\{\mathbf{y}_{J_{1}^{\complement}},\dots,\mathbf{y}_{J_{M}^{\complement}}\right\}$. For each subsampled noisy image $\mathbf{y}_{J_{m}^{\complement}}$, the well-trained Noise2SR takes it as input and generates the corresponding denoising and SR result $f_{\hat{\theta}}\left(\mathbf{y}_{J_{m}^{\complement}}\right)$. Finally, the approximate MMSE estimation of the clean image can be computed below:

where $J_{m}^{c}$ denotes the sub-sampling partition at $m$ time.

As shown in Fig. 3, the network of Noise2SR $f_{\theta}$ consists of a U-Net based Encoder and a Super-resolved Decoder.

Simulated TEM Datasets

Table I shows the quantitative results. Compared with other zero-shot self-supervised learning denoising methods, Noise2SR achieves the best denoising performance. Compared with dataset-based self-supervised denoising methods and supervised denoising, Noise2SR achieves comparable performance.

Fig. 4 shows the qualitative results on two test samples of the simulated noisy \cePt/CeO2 catalyst TEM images. Compared with zero-shot denoising methods, Noise2SR precisely restores the nanoparticle structures and achieves a cleaner background. Moreover, Noise2SR exhibits fewer additional structural patterns in the error maps compared to dataset-based self-supervised denoising methods and supervised denoising methods. In Fig. 5, we compare the intensity profiles on the surface atomic columns for the denoised data obtained from FBI-D and Noise2SR, along with the ground truth. The intensity profiles of both denoised data show a similar overall trend to the ground truth. Compared to FBI-D, Noise2SR exhibits a closer resemblance to the ground truth in terms of both peaks and troughs, with fewer fluctuations.

In this paper, we propose the Random Sub-sampler, a method for generating sub-sampled noisy images from full-size noisy images. To evaluate the effectiveness of randomness in sub-sampler on the denoising performance of Noise2SR, we compare the fix-location and random location sampling strategies within sub-sampler. Fig. 8 illustrates the difference of fix and random location sub-sampling strategies of sampling stride $s=2$. Table II and Fig. 9 show the quantitative and qualitative results, respectively. The results demonstrate that random sub-sampling improves the denoising performance compared to fixed location sub-sampling. Both the denoised results from an individual sub-sampled image $(M=1)$ and MMSE denoised results, random sub-sampling shows the superiority of fixed sub-sampling. In the qualitative result, it is evident that the denoised images obtained from random sub-sampling exhibit clearer backgrounds and sharper atom contours compared to those obtained from fixed location sub-sampling. This observation is further supported by the error maps of the denoising results.