Introduction

Accurate segmentation of brain lesions are essential for treatment planning of brain disorders¹. Unsupervised learning methods, focusing on learning normative distributions from images of healthy subjects, are less dependent on lesion-labeled data, thus exhibiting better generalization capabilities^2,3. A fundamental challenge in learning normative distributions of images lies in the high dimensionality if each pixel is treated as a random variable. Recent deep-learning based methods, such as Variational Auto-Encoders (VAEs) and Generative Adversarial Nets (GANs), have generated very encouraging results^2,4. In this study, we further extended those methods using a subspace-assisted deep generative model to capture both posterior normal and lesion distributions. The proposed method achieved significantly improved segmentation performance across multiple public datasets with stroke, tumor, and multiple sclerosis lesions, in comparison with the state-of-the-art methods.

Methods

The overall framework of the proposed method is shown in Fig. 1, which contains three integral components: 1) posterior normative distribution learned by an image subspace-assisted deep generative model; 2) posterior lesion distribution learned by a tissue subspace-assisted Bayesian classifier; 3) a fusion network to generate the final segmentation.

Problem formulation $$$\\$$$

Given an image with lesion $$$I_t(x)$$$, we represent it as $$$I_t(x)=I_N(x)+\Delta(x),$$$ where $$$I_N(x)$$$ represents an “hypothetical” lesion-free image and $$$\Delta(x)$$$ contains the lesion features. The following posterior distributions are estimated in our proposed method to capture the statistics of $$$I_N(x)$$$ and $$$\Delta(x)$$$ for improved segmentation: a) the position-dependent posterior normal intensity distribution, $$$p(I_N(x_i)|I_t(x))$$$ , at each pixel $$$i$$$; and b) the posterior lesion intensity distribution, $$$p(\Delta(x)|I_t(x))$$$.

Learning $$$\,p(I_N(x_i)|I_t(x))\\$$$
We proposed an image subspace-assisted conditional generative model to estimate $$$p(I_N(x_i)|I_t(x))$$$. After training, $$$p(I_N(x_i)|I_t(x))$$$ can be estimated based on the conditional samples from the learned generative model.
$$$\quad$$$In particular, we first obtained one posterior normal image $$$\hat{I}_N(x)$$$ as the condition of the generative model, leveraging a probabilistic image subspace model constrained by the pre-estimated neighboring normal pixels in $$$I_t(x)$$$ (denoted as $$$W(x)$$$). Specifically, $$$\hat{I}_N(x)$$$ was obtained by: $$\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\hat{I}_N(x)=(1-W(x))\odot\sum_{r=1}^Ra_r\phi_r(x)+W(x)\odot I_t(x), \qquad\qquad\qquad\qquad\qquad\qquad(1)$$where $$$W(x)$$$ is estimated leveraging prior intensity distributions of the normal training images $$$\{\rho_m(x)\}$$$ in a multiscale sense⁵, $$$\{\phi_r(x)\}$$$ is the pre-learned basis functions obtained by applying principal component analysis to $$$\{\rho_m(x)\}$$$, $$$a_r\sim p(\{a_r\})$$$ is the corresponding coefficients obtained by maximum-a-posteriori estimation: $$\qquad\qquad\qquad\qquad\qquad\qquad{\hat{a}_r}=\arg\min\limits_{\{a_r\}}\left\|W(x)\odot\left[I_t(x)-\sum_{r=1}^R a_r\phi_r(x)\right]\right\|_2-\lambda_{\text{coeff}} \log p(\{a_r\}). \qquad\qquad\qquad\qquad\qquad(2)$$$$$\quad$$$After $$$\hat{I}_N(x)$$$ was determined, we learned a conditional GAN to capture the nonlinear mapping from noise samples $$$z$$$ to a real image $$$I_N(x)$$$, conditioned on $$$\hat{I}_N(x)$$$. We adapted the Wasserstein GAN⁶ with improved loss function: $$L_{\text{cWGAN}}(G,D)=\mathbb{E}_{\hat{I}_N,z,I_N}\left[D(I_N)-D(G(z,\hat{I}_N))\right]+\lambda_{\text{GP}}L_{\text{GP}}(D)\\ \qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\quad+\lambda_1L_{\text{L1}}(G,\hat{I}_N)+\lambda_2L_{\text{CE}}(G,\hat{I}_N)+\lambda_3L_{\text{VGG}}(G,\hat{I}_N), \qquad\qquad\qquad\qquad\;(3)$$where $$$L_{\text{GP}}$$$, $$$L_{\text{L1}}$$$, $$$L_{\text{CE}}$$$, $$$L_{\text{VGG}}$$$ is the gradient penalty, L1-loss, pixel-wise cross entropy loss and VGG loss, respectively. The patient-specific posterior normal images with high-fidelity details could be obtained by Monte Carlo sampling conditioned on $$$\hat{I}_N(x)$$$. $$$p(I_N(x_i)|I_t(x))$$$ could then be estimated from the sampled posterior images using a Mixture of Gaussian (MOG) model.

Estimation of $$$\,p(\Delta(x)|I_t(x))\\$$$
To estimate $$$p(\Delta(x)|I_t(x))$$$, we first obtained a global normative spatial-intensity probability distribution: $$\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad p(\hat{I}_N(x_i),c(x_i);\pmb{\theta})=\sum_{k=1}^{3}p(\hat{I}_N(x_i)|c(x_i);\pmb{\theta}_k)\cdot p(c(x_i)), \qquad\qquad\qquad\qquad\qquad\quad\,(4)$$where $$$c$$$ represents the tissue class in $$$\hat{I}_N(x_i)$$$ and $$$p(\hat{I}_N(x_i)|c(x_i);\pmb{\theta}_k)$$$ is the likelihood of the observed intensity $$$p(\hat{I}_N(x_i))$$$ to be tissue class $$$k$$$. To obtain the posterior $$$p(c(x_i))$$$ of $$$p(\hat{I}_N(x_i))$$$, we incorporated a tissue subspace model into the Bayesian classifier for robust segmentation: $$\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\pi_k(c(x_i))=\sum_{r=1}^{R'}\beta_{r,k}\psi_{r,k}(c(x_i)), \qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\quad\;(5)$$where $$$\{\psi_{r,k}(c)\}$$$ are the basis functions of tissue $$$k$$$ obtained on the tissue segmentation labels $$$\{c_1,c_2,\ldots,c_M\}$$$ of the training images. After $$$p(c(x_i))$$$ was obtained, $$$p(\hat{I}_N(x_i),c(x_i);\pmb{\theta})$$$ was fitted using the collection of pixels based on the predicted $$$c(x_i)$$$. With $$$p(\hat{I}_N(x_i),c(x_i);\pmb{\theta})$$$ estimated, the lesion class in $$$I_t(x)$$$ was obtained leveraging voxel-wise Bayesian hypothesis testing, from which $$$p(\Delta(x)|I_t(x))$$$ was estimated using an MOG model.

Fusion network$$$\\$$$

The final classification was determined by fusing the features derived from the learned posterior distributions using a patch-based convolutional neural network, including: 1) position-dependent posterior distributions; 2) multiscale region-level normal/lesion likelihood ratio; 3) multiscale posterior lesion distribution.

Results

To obtain the normal image prior, we used 20000 T1w and 20000 FLAIR brain images from Biobank dataset⁷. The performance was evaluated on simulation data, BraTS2019 tumor (N=288)⁸, MSSEG2016 MS (N=53)⁹ and ATLAS stroke (N=236)¹⁰ datasets, respectively. Fig. 2 demonstrates our proposed method consistently and accurately segmented lesions across various sizes, locations and intensities. As shown in Figs 3-5, our proposed method outperformed the state-of-the-art deep learning-based abnormality detection methods^2,11–15 for challenging lesion/normal interface (tumor), small and isolated lesions (MS) and atrophy brains (stroke). Quantitative analysis also indicated that the proposed method had improved performance compared with the state-of-the-art methods.

Conclusions

This paper presents a new method for lesion segmentation leveraging posterior distributions learned using a subspace-assisted generative model. The proposed method showed significantly improved accuracy and robustness across multiple public brain MR image datasets with various lesion types.

Acknowledgements

This work was supported by Shanghai Pilot Program for Basic Research—Shanghai Jiao Tong University (21TQ1400203), the Program for Professor of Special Appointment (Eastern Scholar) at Shanghai Institutions of Higher Learning, Key Program of Multidisciplinary Cross Research Foundation of Shanghai Jiao Tong University (YG2021ZD28, YG2023ZD22), National Natural Science Foundation of China (62001293).