Generative Tomography Reconstruction
Abstract
We propose an end-to-end differentiable architecture for tomography reconstruction that directly maps a noisy sinogram into a denoised reconstruction.
Compared to existing approaches our end-to-end architecture produces more accurate reconstructions while using less parameters and time.
We also propose a generative model that, given a noisy sinogram, can sample realistic reconstructions.
This generative model can be used as prior inside an iterative process that, by taking into consideration the physical model, can reduce artifacts and errors in the reconstructions
Introduction
This work tackles a prototypical linear inverse problem: reconstruction of Computed Tomography (CT) images in the presence of large amount of noise.
The goal of CT reconstruction is to infer an image x , given a noisy sinogram y
δ ∼ n ( Ax ) , where A is the discrete Radon transform [3] and n ( · ) is the noise distribution.
Solving such a reconstruction problem is difﬁcult due to the ill-conditioning of Radon transforms.
It is also a problem of high practical relevance and impact, given the widespread use of CT algorithms in applications such as medical imaging [12, 22, 2], luggage scanning [11] and damage detection in concrete structures [7].
This work is centred around two original contributions.
First, we design an end-to-end differentiable architecture that directly maps a noisy sinogram to a denoised reconstruction.
By combining reconstruction and denoising functionalities into a single model, we are able to improve reconstruction quality while reducing inference time.
Existing ML approaches to noisy CT reconstruction [8, 21, 10] use a convolutional neural network (CNN) to denoise a reconstruction computed by a non-adaptive algorithm.
Contrary to that, we apply the learned reconstruction and denoising process directly to the sinogram.
This allows the model to reduce noise before it is ampliﬁed by the reconstruction process.
The second contribution leverages the differentiable architecture to propose a generative model that can sample reconstructions which are realistic and compatible with sensor readings.
The reconstructions produced by the proposed generative model can be iteratively improved by taking into consideration the physical model of sinogram generation.
This iterative algorithm can reduce reconstruction errors and biases, guaranteeing that the produced solution is compatible with what is measured by the sensors.
GAN-based approaches in literature for solving linear inverse problems [20, 4, 15] compute reconstructions iteratively by improving an initial random guess.
In contrast to that, our approach can generate good solutions even without iterating, making the iterative improvement process much faster.
As a ﬁnal note, we point out the generality of our approaches, which allow a straightforward application to other linear inverse problems, such as limited/sparse angle tomography, image super-resolution and image deblurring.
Sinogram-Based Tomography Reconstruction
Before delving into the details of the two contributions, we brieﬂy summarize the noise model used to simulate sensor readings.
We take into consideration shot and electronic noise, and quantization errors as follows
[Equation]
where exp( s ) is the X-ray intensity, b is the number of bits used by the detector and k is a scaling parameter.
We assume that, given readings r (that are integer values), sinograms are distributed as y | r ∼ N ( µ ( r ) , diag ( σ 2 ( r ))) , where µ ( · ) and σ ( · ) can be modelled by CNNs.
We train µ ( · ) and σ ( · ) networks to minimize the negative log-likelihood (refer to Appendix A for more details).
2.1 Reconstruction Model
Inspired by ﬁltered back-projection (FBP) [9], which reconstructs x as (cid:101) x
= A T f ( y ) , where f ( · ) is a ﬁxed linear ﬁltration, we design our reconstruction model as g ( r ) =
g
( A T g 1 ( r )) where both g 1 and g 2 are trainable neural networks.
The Radon transform is ill-conditioned and FBP is wellknown to greatly amplify noise.
Our rationale is that by combining reconstruction and denoising, we allow the model to mitigate the reconstruction errors caused by the ill-conditioning of the Radon transform.
Furthermore, since the output of g 1 ( r ) is not limited to be a single channel, it allows richer information ﬂow from the sinogram to the image.
Brieﬂy, our reconstruction model is composed of two U-nets [19] connected by Radon backprojection (multiplication by A T ).
The ﬁrst one ( g 1 ) processes the sensor readings and outputs a multi-channel sinogram, each channel is transformed into an image using Radon backprojection and fed to the second network.
The ﬁrst layer of g 2 concatenates to its input (of shape ( ch , 256 , 256) )
a 4 dimensional positional embedding (of shape (4 , 256 , 256) ) which is used by the network to encode pixel coordinates.
The down-sampling blocks of the U-net are composed of a strided convolution followed by a sequence of residual blocks [6].
Each U-net up-sampling block takes the output of the previous block, double its resolution using bilinear up-sampling, and concatenates it with the skip connection coming from the down-sampling block at the corresponding resolution.
Following, the concatenation channels are weighted using channel attention [14] and inputted to a 1 × 1 convolution that shrinks their number.
Finally they are processed by a sequence of residual blocks.
A detailed account of the architecture is provided in Appendix B.
The key point in our convolutional scheme is that the Radon backprojection between the two U-nets is differentiable.
Therefore the model discussed above can be trained end-to-end by jointly optimizing g 1 and g 2 to minimize an image reconstructions loss.
2.2 Generative Model
We design an end-to-end generator model that, given noisy sensor readings r , can sample different possible reconstructions.
We show how, by combining it with a discriminator and the noise model, it can be used to iteratively improve reconstructions.
We build on the Wasserstein GAN (WGAN) [1] framework and train a generator G to minimize
E r,z (cid:104) (cid:107) AG ( r, z ) − µ ( r ) (cid:107) 2 σ ( r ) + λD ( G ( r, z ))
(cid:105) where (cid:107) v (cid:107) 2 σ ( r ) def = (cid:88)
i (cid:18) v i σ i ( r ) (cid:19)(1) where z is uniformly distributed over the sphere and λ is a ﬁxed regularization parameter.
Simultaneously, we train a discriminator D to minimize
E x
[ D ( x )] − E r,z [ D ( G ( r, z ))] subject to (cid:107) D (cid:107) L ≤ c where the Lipschitz-norm constraint on D is enforced using Spectral Normalization [16].
The generator G ( r, z ) is an end-to-end model following the architecture described in Section 2.1.
To make the computation of ∇ z G ( r, z ) efﬁcient, we feed z to the third last block of the g 2 network, so that back-propagating the gradient to z only requires to go through few layers.
The structure of this modiﬁed U-net block is depicted in Figure 1 to highlight the additional input from z .
The ﬁrst layer of this block concatenates its inputs (by applying bilinear up-sampling to the ones coming from the previous block) and applies a 3 × 3 convolution followed by a PReLu activation [5]. The effect of z on the output should be modulated depending on the amount of noise and uncertainty in the reconstruction. Therefore the last 32 channels produced by the ﬁrst layer of the block are fed to an
Conv Split exp ×
Upscale z
Concat
Res Blocks g
[Table]
SSIM
Parameters
FPS
[Table] [Table]
XL
[Table]
[Table]
XXL
[Table]
XXS
[Table] [Table]
XS
M [Table] [Table]
L
[Table]
exponential activation and used to scale the elements of z (bottom part in Figure 1).
The remaining channels are then concatenated with the scaled version of z and fed to a sequence of residual blocks.
The parameter z has shape (32 , 16 , 16) and it is up-scaled to (32 , 64 , 64) using bilinear interpolation.
If we train the generator using the loss described in (1) it learns to ignore z .
To avoid this problem, we regularize the output of G by pushing the distribution of AG ( r, z ) towards N ( µ ( r ) , diag ( σ 2 ( r ))) .
During training we generate two different reconstructions for each noisy reading.
Then, we ﬁt a normal distribution with diagonal covariance on the two corresponding sinograms and compare it against N ( µ ( r ) , diag ( σ 2 ( r ))) using the Kullback–Leibler divergence
[Equation]
2 log (cid:18) σ ( r ) σ p (cid:19) + (cid:18) σ p σ ( r ) (cid:19) 2 (cid:35) where σ 2 p =
( y 1 − y 2 ) .
Using this loss in place of the ﬁrst term of (1) solves the mode collapse problem without reducing training stability.
As a ﬁnal note, the reconstruction produced by our generator can be iteratively improved using projected gradient to minimize (cid:107) AG ( r, z ) − µ ( r ) (cid:107) 2 σ ( r ) + λD ( G ( r, z )) subject to (cid:107) z (cid:107) = 1 .
Experimental Assessment
We use images from the DeepLesion dataset 1 , cropped to guarantee that pixels outside of the inscribed circle are zero and rescaled to 256 × 256 (see Appendix D for more details about preprocessing).
We split train/validation/test by patient identiﬁer and use multiple signal intensities (therefore noise levels) during training and testing.
No augmentation is used, sensor readings are not ﬁxed but randomly simulated before each training iteration.
We compare the reconstruction quality obtained by our end-to-end approach and an FBP model, also varying the dimensions of the U-net components g 1 and g 2 .
We use conventional names (XXS, XS,
S, S-64, M, L, XL, XXL) to refer to increasing U-net dimensions (refer to Appendix B for further details).
All the models are trained for 10 epochs using the RAdam [13] optimizer, batch size ofand mixed precision training.
The learning rate is exponentially increased from 0 to 3 · 10 − 4 during the ﬁrst 5000 batches, then halved every 80000 batches.
We measure the quality of reconstructions using structured similarity (SSIM) and compute inference speed as the number of frames per second (FPS) obtained with the largest possible batch size on a Tesla V100 GPU.
Results are shown in Table 2, from which it can be noticed that end-to-end models achieve a more accurate reconstruction than FBP based models at comparable number of
FPS.
End-to-end models scale better, a +0 .
7% improvement in SSIM requires to increase the number of parameters 2 .
6 × , compared to FBP models where an improvement of +0 .
3% costs 2 .
5 × more parameters.
Furthermore end-to-end models are much more parameter-efﬁcient, given that their largest model is only slightly larger than the smallest FBP model.
Figure 4 depicts the progress of the projected gradient algorithm on two candidate reconstructions using a learning rate of 10 − 4 .
Notice in particular how the reconstruction of the spinal cord is improved by the iterative algorithm.
Conclusions and Future Work
We proposed an end-to-end architecture for tomography reconstruction characterized by higher quality and performance than FBP-based denoising models.
The key point of our approach is the incorporation of the Radon transform as a differentiable component of the neural architecture, allowing to build the reconstruction directly from the sinogram.
The end-to-end architecture can be used to design generative models capable of sampling multiple possible reconstructions.
Also, we show that by taking into consideration the physical model, it is possible to iteratively improve the reconstructions computed by the generative model.
Our work can be easily extended to sparse angle and limited angle tomography or to other linear inverse problems.
A Noise Model
Without taking into consideration noise the ideal sensor readings would be exp( s − y ) , where exp( s ) is the intensity of the X-ray emitted by the machine.
In practice due to shot noise and electronic noise readings can be modelled as
[Equation]
In this setting noise depends on the signal intensity (the stronger the signal the weaker the noise) and on the variance of the normal distribution that emulates electronic noise.
Notice that, because of the
Poisson distribution, noise will not be uniform across the sinogram and will be higher where y is large.
In practice the analog value at each pixel of the sensor is converted into a digital integer value, this process creates quantization error.
If readings are stored as b bits values then this conversion process can be modelled as
[Equation]
where k is a normalization value that depends on the scale of the measured signal.
Given sensor readings r (a matrix of integers in range [0 , 2 b ) ) we would like to estimate the posterior probability over noiseless sinograms π ( y | r ) .
We make the simplifying assumption of y | r ∼ N ( µ ( r ) , diag ( σ 2 ( r ))) , where µ ( · ) and σ ( · ) are CNNs.
Because both the mean and variance are modelled by CNNs intra-pixel relationships can be modelled.
We ﬁrst train µ ( · ) to minimize
[Equation]
and, after µ ( · ) is trained, we use it to train σ ( · ) to minimize the negative log likelihood
[Equation]
Finally we use knowledge distillation to combine µ ( · ) and σ ( · ) into a single CNN.
The architecture of the end-to-end reconstruction model is depicted in Figure 5. Both g 1 and gare instances of the U-net architecture [19]. Blocks that have more than one input, ﬁrst up-sample each input to the largest resolution using bilinear up-sampling, then concatenate all the inputs on the ggr
Embedding
U-net Block
U-net Block . . .
U-net Block U-net Block
Conv 2D
A T
Pos Features
U-net Block
U-net Block . . .
U-net Block U-net Block
Conv 2D g
[Figure]
gr
Embedding
U
[Figure]
Block
[Figure]
Block . . .
U
[Figure]
Block
U-net Block
Conv 2D
A T
Pos Features
U-net Block
U-net Block . . .
Modiﬁed Block z
U-net Block
Conv 2D channel dimension.
The architecture of the proposed generator is shown in Figure 6 and is analogous to the end-to-end reconstruction model.
The only difference is the modiﬁed block that is used to receive z (bottom right of Figure 6) which is described in Section 2.2.
The down-sampling blocks of the U-net are composed of a strided convolution followed by a sequence of residual blocks [6].
Each U-net up-sampling block takes the output of the previous block, double its resolution using bilinear up-sampling, and concatenates it with the skip connection coming from the down-sampling block at the corresponding resolution.
Following, the concatenation channels are weighted using channel attention [14] and inputted to a 1 × 1 convolution that shrinks their number.
Finally they are processed by a sequence of residual blocks.
We use conventional names (XXS, XS, S, S-64, M, L, XL, XXL) to refer to increasing U-net dimensions, the only difference between S and S-64 is the number of channels.
The number of channels used by the residual blocks is doubled every two strided convolutions.
The ﬁrst block of
XXS, XS, S uses 32 channels, while the ﬁrst block of S-64, M, L, XL, XXL uses 64 channels.
The number of residual blocks used inside each U-net block is shown in the following table:
In end-to-end models the latest convolution of g 1 reduces the number of channels before Radon backprojection.
The smallest end-to-end model (which uses a XXS U-net for g 1 and a S-64 U-net for g 2 ) uses 16 channels for Radon backprojection, while the other two end-to-end models usechannels.
Resolution
XXS
XS
S
M
L
XL
XXL                         -       - -      - - -     - - -     - -      -
C Motivation for the Generator Loss
A standard approach for solving inverse problems is to use regularization, this amounts to solving
[Equation]
where D ( x ) is a penalty term that promotes desired properties of x and λ is the regularization parameter.
First, we modify this formulation to take into consideration our sinogram estimation:
[Equation]
where
[Equation]
With this formulation we are taking into consideration the non-uniform conﬁdence in the estimated sinogram values.
There is still the problem that we are directly optimizing over images, and we still have to deﬁne D ( x ) .
Instead of optimizing x we introduce a generator G ( r, z ) and optimize over the free parameter z that lies on the m + 1 dimensional sphere:
[Equation]
D Dataset Preprocessing
We use the DeepLesion dataset (https://nihcc.app.box.com/v/DeepLesion) which is a large dataset of annotated tomographic images stored with 16 bit precision.
We only use the images and don’t consider the annotations for training our model.
Images are preprocessed as follows:
Subtract 32768 from the pixel intensity to obtain the original Hounsﬁeld unit (HU) values.
HU is a measure of radiodensity where air has density − 1000 HU and water 0 HU.
Create a binary version of the image by thresholding pixel with value larger than − 9050 HU (threshold is a bit higher than the density of air).
Find the smallest circle that contains all the pixels of the binarized image and to crop the corresponding square region from the image.
If the cropped region is smaller than 256 × 256 discard it, otherwise rescale to 256 ×and add to the dataset.
E Implementation Details
Training the model requires computation and differentiation of Radon forward and backward projections (corresponding respectively to products with A and A T ).
These operations are linear and therefore differentiable, we use the implementation offered by the TorchRadon library [18] which is integrated with PyTorch [17] and allows the backpropagation of gradients.
F Additional Generated Reconstructions
The following ﬁgures contain additional images generated by our model using all the tested signal levels.