Combination of edge enhancement and cold diffusion model for low dose CT image denoising

Denoising diffusion probabilistic models (DDPM)

DDPM refers to a generation model that learns target data distribution from samples [22]. DDPM consists of two Markov processes: a fixed forward process and a learning-based reverse process, as shown in Figure 1.

Figure 1:

Diagram of denoising diffusion probabilistic models. The Markov chain of forward (reverse) diffusion process of generating a sample by slowly adding (removing) noise.

Forward process: Starting from a sample x₀∼q(x₀) of a clean image, Gaussian noise is gradually added according to the following probability:

Type of N ( · ) said the gaussian probability density function, β_1:T complying with β _t ∈(0,1) all of the t=1, ⋯, T. The latent variable x_1:T has the same dimension as the original image sample x 0 ∈ R n , and when T is large enough and an appropriate β _t scheduling is selected, the latent variable x _T approaches a Gaussian distribution. By changing the α _t :=1−β _t and α ‾ t = ∏ s = 1 t α s parameters, x _t can be written as a linear combination of the noise ϵ and x₀:

where ϵ ∼ N ( 0 , Ι ) . This allows for a closed expression for the marginal distribution of x _t samples given x₀:

Reverse process: Since the inverse procedure q(x_t−1|x _t ) depends on the entire data distribution and is difficult to handle, a neural network can be used to learn the parameterized Gaussian probability p _θ (x_t−1|x _t ), as follows:

where z ∼ N ( 0 , Ι ) . Since the model learns a reverse Markov chain running in inverse time from x _T to x₀, we call it a reverse process to estimate a clean image x₀ from a partially noisy image x _T .

In the reverse process, x _t in equation (5) can be replaced by (2) to obtain the relation of x_t-1 represented by x₀ and noise, which can also be interpreted as that each step of reverse sampling actually predicts the initial image. Then the noise map in the current state is linearly superimposed on the predicted image, resulting in a new image that serves as the image for the next iteration.

Cold diffusion

The cold diffusion model is a generalized diffusion model that extends the diffusion and sampling of Gaussian noise to any type of degradation, such as adding various types of noise, blurring, downsampling, etc. Specifically, given an image x₀ from the training data distribution Q, the image x₀ (cold state) is gradually degraded to an image x _T sampled from a random initial distribution P (hot state) using a specific degradation operator D(·), such as a Gaussian distribution where T is the total number of time steps of diffusion. Then, the image x _t of any time step t in the diffusion process is defined as x _t =D(x₀,x _t ,t), where t corresponds to the degree of degradation, and for any t, the operator D(·) should be continuous.

In LDCT image denoising, the NDCT image is taken as the starting point x₀ and the LDCT image is used as the end point of the diffusion process which is x _T . The degradation operator D(·) is introduced to ensure that the middle image x _t of the diffusion process maintains the same expectation as x₀, without introducing additional CT number shift, which is defined as follows [25]:

where x _T is a random noise with a known distribution and α _t <α_t-1, ∀1≤t≤T.

In the reverse process, x _T is first sampled from P, and then the diffusion process is reversed using the recovery operator R(·), which can be expressed as:

In practice, R(·) is a neural network parameterized to θ, which can be optimized by the following objective functions:

It is worth noting that for any t, R _θ (·) can directly generate the recovered image x ˆ 0 from x _T . However, such a one-step prediction may produce a blurry image x ˆ 0 and serious detail loss.

In order to solve this problem, the cold diffusion model follows the annealing sampling algorithm in the classical diffusion model, adopts the “restoration-redegradation” sampling algorithm, and gradually generates the image through T sampling steps. From the predicted value x ˆ 0 , the image x ˆ t − 1 of time step t−1 can be calculated as follows:

Although this iterative sampling algorithm can produce a sharper image than the single step prediction, the prediction error between x₀ and x ˆ 0 may cause the error between the true value x_t−1 and x ˆ t − 1 of the time step t−1 to increase. Errors accumulate gradually during the sampling process, and the prediction bias of R _θ (·) may be further worsened. Therefore, Bansal et al. [26] proposed an improved sampling algorithm to reduce this cumulative error:

where x ˆ T = ( x t − α t x ˆ 0 ) ∕ 1 − α t . While the improved sampling algorithm in equation (10) mitigated the problem of error accumulation and was shown to produce better image quality, it did not correct the misalignment between the input and its corresponding time step, which could result in non-negligible shifts in pixel values. Therefore, Timestep adjustment module is introduced into the network in this paper to reduce the cumulative error caused by imperfect recovery network and misalignment in cold diffusion. We will introduce the details in the part of network architecture.

Convolutional block attention module

As shown in the lower left corner of Figure 2, CBAM [27] is a convolutional block attention module that combines channel and spatial. Given an input feature, CBAM can derive the attention graph along two independent dimensions, channel dimension and space dimension respectively, and then multiply the attention graph with the input feature for adaptive feature thinning.

Figure 2:

Overall architecture of CECDM network. The lower left corner is the CBAM module, and the lower right corner is timestep adjustment module.

In terms of channel attention (the blue part in Figure 2), the spatial dimension of the input image graph was compressed, the channel attention was calculated, and then aggregated. Two descriptors were obtained by means of average pooling and maximum pooling. These descriptors are forwarded to a three-layer shared multilayer perceptron (MLP) to generate a attention map. The inputs for each MLP are then summed by unit and the channel attention is calculated using the sigmoid function.

Spatial feature relationships are then used to supplement channel attention (green in Figure 2). Spatial attention is focused on the information part by applying average pooling and maximum pooling channel calculations and then concatenating the two to get a single feature descriptor. In addition, convolutional layers are used on the connected feature descriptors to generate spatial attention graphs, which are encoded to emphasize or suppress. The channel information of feature graph is aggregated by average pooling feature and maximum pooling feature, and then convolution and convolution are performed to generate two-dimensional spatial attention maps.

Network architecture

The overall structure of the CECDM network proposed in this paper is shown in Figure 2. Each timestep in the training process is divided into two stages, the left is the first stage in the training process, the right is the second stage, the two parts of the network composition is the same. In the first stage, the input image x _t estimates x₀ through the recovery network R _θ (·), which is equivalent to one-step prediction of x₀. In the second stage, x ˆ 0 obtained in the previous step is redegenerated, and x ˆ t − 1 is calculated by equation (9) as the input of the network, and then the same network R _θ (·) is used for processing, and the output x ˆ ˆ 0 is the input x_t−1 of the next time step. This part is achieved through the Timestep Adjustment Module, as shown in the blue dashed box in the lower right part of Figure 2.

The whole recovery network R _θ (·) adopts UNet network framework. The input images are first processed by 3 × 3 sobel convolution module, then average pooling undersampling and TAM module in corresponding stage, and finally output by a combination of 3 × 3 Conv and ReLU, which is introduced into CBAM module before upsampling. The CA module uses average pooling and maximum pooling to compress the spatial dimension of the feature map and retain more texture information of LDCT image. The SA module supplements the CA mechanism and reduces the number of channels to one through the dimensionality reduction filtering of 7 × 7 convolution kernel. Finally, the SA feature map is obtained through the sigmoid activation function. Enrich the input information up-sampled. The feature is modulated by TAM module after each up and down sampling.

Timestep adjustment module

The timestep embedding feature Cond _t is obtained by sinusoidal position coding t and MLP. The output of the first stage is obtained by adding the module input and t time feature:

The timestep input of the second stage is t−1, and x ˆ 0 and x _T are spliced in channel dimension. Two modulation factors γ and β are obtained by convolution RELU and other operations F ξ . γ is multiplied by time feature and added with β and input to get the final output f_t−1.

Sobel convolution block

In order to enrich the input information of the model, a trainable sobel convolution block is used as the edge enhancement module of the network, and some changes are made to it while retaining the excellent edge extraction ability of the sobel operator, as shown in Figure 3.

Figure 3:

Learnable sobel operator.

When processing LDCT images, it is found that the edge information extracted by using a single traditional sobel operator is incomplete. In this case, the absence of contour information is more serious. Different from the traditional fixed-valued Sobel operator [28], EDCNN [15] proposes a trainable Sobel operator, which defines four operators in the horizontal, vertical and diagonal directions as a group, and this module is also used in DEformer [29], so that the parameter values can be adaptively adjusted during the optimization training process to extract the edge information with different intensities. In order to make the CECDM model extract edge information more accurately, we make some modifications on this basis by changing the learnable parameter α in the sobel algorithm from 0 to 0.1, and other parameters remain unchanged.

Loss function

In this paper, we take the composite loss of MSE+Charbonnier loss as the final training target of the network. The output of the first stage of the network is R _θ (x _t ,t), and the output of the second stage t−1 is x ˆ ˆ 0 . The MSE loss of the two stages is calculated respectively to obtain L_MSE.

Charbonnier Loss [30], 31] is an approximation of L1 loss and is used to improve the performance of the model. L_char calculates the loss between the network output x ˆ ˆ 0 and x₀ with parameter ε=1 × 10⁻³.

The corresponding formula for the final compound loss L C o m p o u n d _ l o s s is as follows:

where λ₁ and λ₂ are predefined constants, λ₁=1.5, λ₂=0.05 in this paper.

Mayo dataset

The 2016 NIH-AAMP-Mayo Clinic low-dose CT Grand Challenge [32] provided a total of 2,378,512 × 512 abdominal CT scans from 10 patients. Under the parameter setting of 120 kVP tube voltage and 200 mAs effective dose, the image was scanned and reconstructed with CT scanner, and the image section thickness was 1 mm and 3 mm. To simulate images at low dose levels, Poisson noise was inserted into the projected data for each case to obtain LDCT images corresponding to a noise level of one-quarter of the full dose. We used 3 mm thickness slices, 2,167 images of nine patients for training, and 211 images of one patient for testing.

Piglet dataset

Real Piglet dataset [33] provided 850 CT images of 512 × 512 piglets with section thickness of 0.625 mm. Under the condition of 100 kVP tube voltage, GE scanner adjusted the tube current of 300 ∼ 15 mAs to perform CT scanning and reconstruction on piglets. Among them, 300 mAs is used as the conventional dose level, and the remaining dose level images are obtained by reducing the dose to 50 %, 25 %, 10 % and 5 % of the original dose, respectively. Four dose data of 50 %, 25 %, 10 % and 5 % were selected for training and testing. 525 Piglet images were selected for each dose as training input and 40 images were selected for testing.

Experimental detail setting

We use U-Net as the backbone of the proposed CECDM with an input size of 3 × 512 × 512. We optimized it using the Adam optimizer with a batchsize of four and a learning rate of 2 × 10⁻⁴, for a total of 100 k iterations of training. α₁, … , α _T is set to the cosine timesteps schedule [34] from 0.999 to 0, where s=0.008. The experiment was based on the Pytorch framework and trained and tested the network using a computer configured as an NVIDIA GeForce RTX 2080 SUPER GPU.

Analysis of experimental results

In order to test the effectiveness of the proposed CECDM algorithm, the proposed algorithm is compared with REDCNN [12], QAE [13], EDCNN [14], CTformer [19] and CoreDiff [25] in terms of visual assessment and quantitative assessment. Among them, objective evaluation indicators in quantitative assessment include the peak signal-to-noise ratio (PSNR) based on the grey diference of the pixel, the structural similarity measure (SSIM) based on the structural diference [35], the gradient similarity deviation (GMSD) based on the change of the gradient [36], the feature similarity measure (FSIM) based on the diference of the feature [37], and the implementation of the multiscale pixel domain (VIFs) based on visual perception [38] are used. GMSD represents the pixel similarity of the gradient amplitude map between the reference image and the processed image, and the lower the GMSD score of the processed result is. In addition to GMSD, the higher the value of other indicators, the better the quality of CT images after noise removal. The comparison experiments are derived from the official code, and their parameters are set according to the recommendations of the paper.

Mayo dataset visual assessment

Figure 4 show the processed results of two representative slices (denoted as Case 1 and Case 2) from the Mayo test dataset by different methods. All CT images of the axial view are displayed in the [−160, 240] HU window. Due to insufficient photons from incident X-rays, there are a lot of speckle noise and edge artifacts in LDCT images, making it difficult for physicians to determine clinically significant lesion information. From NDCT images (Figure 4(b1) and (b2)), clear lesions can be observed, but there are serious artifacts and noise in LDCT images, as shown in Figure 4(a1) and (a2), so it is difficult to sketch accurately. To show the noise reduction more clearly, we also zoom three ROI(region of interesting) areas (ROI1 in the lower right corner of Case 1, ROI2 in the upper left corner of Case 2, and ROI3 in the lower right corner) and marked them with red arrows and blue circles for better comparison. In general, these six methods can achieve a certain degree of denoising effect, and different denoising methods have different degrees of suppression of noise and artifacts in LDCT images.

Figure 4:

Comparison of processed images by different methods for Case 1 and Case 2. (a1, a2) LDCT, (b1, b2) NDCT, (c1, c2) REDCNN, (d1, d2) QAE, (e1, e2) EDCNN, (f1, f2) CTformer, (g1, g2) CoreDiff, (h1, h2) CECDM. ROI1 in the lower right corner of Case 1, ROI2 in the upper left corner of Case 2, and ROI3 in the lower right corner.

In contrast, REDCNN (Figure 4(c1) and (c2)), while eliminating some noise, suffers from image blurring and images that are too smooth and texture details are lost, such as the details marked by the red arrow in ROI2, the upper left corner of Figure 4(c2). QAE (Figure 4(d1) and (d2)) and CoreDiff (Figure 4(g1) and (g2)) can significantly eliminate noise, and the detail retention effect is better than REDCNN to a certain extent, but the visual effect of these two methods is blurry. EDCNN(Figure 4(e1) and (e2)) and CTformer (Figure 4(f1) and (f2)) are better at preserving image detail than the other methods, but fine striped artifacts are still visible at the red arrows of ROI1 and ROI2. In contrast, CECDM achieves better results in effectively removing noise/artifacts and preserving tissue/structure. The lesion contours of ROI1 (blue circles shown in Figure 4(e1) and (h1)) are easy to identify. It can be found that the lesion contours of EDCNN and CECDM are the clearest, while those of other methods are fuzzy, which confirms the role of edge enhancement module. In addition, CECDM has the clearest tissue edges in ROI3 in Figure 4 (h2). Compared with other images in ROI3 in Figure 4, CECDM has the strongest ability to remove noise and artifacts, while other methods have significant residual artifacts. The results show that CECDM has good texture preservation and noise suppression ability.

Mayo dataset quantitative assessment

In order to quantitatively evaluate the denoising effect of different methods, the average PSNR, SSIM, GMSD, FSIM and VIFs values of all test results of different algorithms are given in Table 1. The best values are shown in red and the next best values are shown in blue. The method in this paper has good quantitative performance on all objective indicators, and the improvement of PSNR and SSIM is also obvious.

As can be seen from Table 1, the difference between REDCNN, QAE and CTformer in PSNR is small. But REDCNN had the lowest GMSD score, consistent with the visual judgment that the image was too smooth and had low contrast. EDCNN performs poorly on PSNR, GMSD, and VIFs, but has a smaller gap with other methods on other quantitative measures. QAE and CTformer are superior to the two typical CNN-based methods above on SSIM and FSIM. CoreDiff and CECDM are ahead of other methods on all indicators, which proves that diffusion models hold encouraging promise in the field of LDCT post-processing.

CECDM has the best value on several indexes, indicating that CECDM can achieve better noise/artifact suppression and feature information retention, and the results are close to the corresponding NDCT images. In summary, CECDM is superior to other comparative methods in visual effect and quantitative index analysis.

Piglet dataset visual assessment

Figure 5 shows the treatment results of representative piglet slices (Case 3) by different methods at different doses (50 %, 25 %, 10 %, 5 %). All CT images are displayed in window [−160, 240] HU. Figure 5(a1)–(a4) shows LDCT images at 150 mAs, 75 mAs, 30 mAs and 15 mAs. Figure 5(b1)–(b4) shows the corresponding NDCT images. Six comparison algorithms RED-CNN, Q-AE, EDCNN, CTformer, CoreDiff and CECDM. Comparing LDCT and NDCT images, it is found that the lower the X-ray dose, the richer the noise information in LDCT images. When the tube current drops to 15 mAs, the quality of LDCT image is significantly lower than that of NDCT image, and it is difficult to observe the tissue structure and detailed information.

Figure 5:

Image comparison of different processing methods in piglet dataset. (a1–a4) LDCT (the tube currents are in order 150 mAs, 75 mAs, 30 mAs and 15 mAs), (b1–b4) NDCT, (c1–c4) REDCNN, (d1–d4) QAE, (e1–e4) EDCNN, (f1–f4) CTformer, (g1–g4) CoreDiff, (h1–h4) CECDM.

We zoom two ROI areas (marked by red rectangles in Figure 5(b)) for better comparison, trained and tested the four doses of LDCT separately, and we can see that above methods have different performance in noise reduction effect when tested on 150 mAs and 75 mAs datasets. Similar to the results from the Mayo dataset, RED-CNN, EDCNN, CTformer, and Q-AE were effective in suppressing noise at high doses (i.e. 50 % and 25 %). However, the RED-CNN image is too smooth (ROI regions (c1), (c2) in the lower right corner of Figure 5). Although EDCNN can retain edge information and has good noise reduction effect, there is still a problem of blurred details in the image (as shown in Figure 5(e1)–(e2)). CTformer can obtain improved images of better quality, but there is a smoothing phenomenon (as shown in Figure 5(f1)–(f2) top left corner ROI regions). The results of CoreDiff and CECDM are superior to those of the above methods, and it can be seen that the method proposed in this paper has good performance in both ROI regions (Figure 5(h1)–(h2)). However, when the tube current is reduced to 30 mAs and 15 mAs respectively, the image quality is seriously polluted. The two diffusion model-based methods proposed by us, CECDM and CoreDiff, have remarkable denoising effects and are obviously superior to other methods.

Piglet dataset quantitative assessment

Figure 5 shows the PSNR and SSIM indicators for the optimal and sub-optimal values of different noise levels, denoted in red and blue respectively. As can be seen from Figure 5, the X-ray dose is closely related to the LDCT reduction results of the six algorithms on the piglet test set. The higher the X-ray dose, the better the image index, and the lower the dose, the worse the image index.

It can be seen that when the tube current is 150 mAs and 75 mAs the indexes of RED-CNN are improved, and the optimal and sub-optimal values are obtained in some indexes. However, the index gap of other methods is not obvious. When the tube currents are 30 mAs and 15 mAs, CoreDiff and CECDM methods based on cold diffusion are obviously better than other methods, and most of the indexes are optimal. In general, the CECDM method proposed in this paper has achieved good results on measurements of LDCT image denoising at different doses, and has strong robustness.

Ablation studies

In this section, we design several sets of ablation studies to analyze and verify the performance of the proposed CECDM model under different network component and compound loss. In addition, it is worth mentioning that the cross-validation strategy is no longer used in this part of the experiment, but a set of data tests from the Mayo dataset are fixed.

Diffusion timesteps T

The performance of CECDM in training and inference was evaluated at T=1, 10, 50 and 200. Figure 6 shows the images after noise reduction with different T values. When T=1, CECDM is a one-step restoration, the edge of the image after noise removal is blurred, and the texture details are not clear. With the increase of T, the organizational boundary becomes clearer. At T=200, the degree of detail retention of the image is better than that of the image at T=50, but the inference time also increases significantly.

Figure 6:

Ablation studies of CECDM with different T settings (a) LDCT, (b) NDCT, (c) T=1, (d) T=10, (e) T=50, (f) T=200.

When T=10, the PSNR and SSIM values of CECDM are the highest, and the image after noise removal is the closest to NDCT visually. When T≥50, the quantitative performance of CECDM decreased gradually. When T=200, PSNR increases to a certain extent. Therefore, considering the sharpness of the denoised image, quantitative performance, and inference time of the method, T=10 is a suitable setting for CECDM.

Network components and compound loss functions

TAM and Sobel+CBAM modules are used to improve the network structure and compound loss function is used to optimize the network performance. In order to prove the effectiveness of these improvements, a series of ablation studies are conducted to evaluate their effects on the performance of the noise reduction network. The results of quantitative analysis are shown in Table 2, and it can be seen that our proposed CECDM method is the most effective. All four different combinations of improvements have been improved to some extent, with Sobel+CBAM showing the most significant improvement. The proposed improvements have been improved to some extent, among which Sobel+CBAM has the most obvious improvement effect. Although the overall index of other improvements has not been greatly improved, there are still some gaps in some small details. In summary, the improvement strategy proposed in this paper has a positive effect on improving the quality of denoised images.