Unsupervised Brain Tumor Segmentation with Image-based Prompts

Since brain tumors generally appear as hyper- or hypo-intensity on MR images, in PL-BTS the image-based prompt is designed to indicate whether there is hyper-/hypo-intensity at a voxel. Equivalently, a prompt model is trained to determine for each voxel of an input image: does the voxel in the input image have hyper-/hypo-intensity? To train such a model, we generate artificial hyper-/hypo-intensity on images without brain tumors with hand-crafted rules, where no manual annotation is needed.

Suppose we are given a set $\mathcal{X}=\{\mathbf{X}_{i}\}_{i=1}^{N}$ of images without brain tumors, where $\mathbf{X}_{i}$ is the $i$-th image and $N$ is the total number of these images. Based on $\mathbf{X}_{i}$, we can generate a synthetic image $\mathbf{X}$ with regions of hyper-/hypo-intensity and obtain the corresponding mask $\mathbf{Y}$ of abnormality indicating the hyper-/hypo-intensity as

Here, $T(\cdot)$ represents an intensity transformation that generates hyper-/hypo-intensity, $\mathbf{A}$ is a weight image with intensities ranging from zero to one, $\odot$ represents voxelwise multiplication, and $a$ is a threshold that determines the mask $\bf{Y}$ based on $\bf{A}$. Eq. (1) mixes the original image without brain tumors and the artificial hyper-/hypo-intensity to generate an image with intensity abnormality, and the threshold $a$ (empirically set to 0.1) is used to account for the partial volume effect in the buffering region where normal tissue and abnormality are mixed. The key to the prompt design is the designs of $T(\cdot)$ and $\mathbf{A}$, which determine the appearance and shape of the hyper-/hypo-intensity, respectively. Note that we do not seek to generate realistic tumors on the images, as it is extremely challenging, if possible at all, without any annotation.

As brain tumors comprise large clusters of abnormally bright signals or dark signals, $T(\cdot)$ is designed to generate hyper-/hypo-intensity with smooth textures. Specifically, $T(\cdot)$ blurs and scales the image to produce hyper-/hypo-intensity as

Here, $\mathbf{Blur}(\mathbf{X}_{i})$ is a 3D Gaussian blurring function applied to $\mathbf{X}_{i}$, and $\lambda$ controls the magnitude of the hyper-/hypo-intensity. The specification of $\lambda$ is introduced later. Since $T(\mathbf{X}_{i})\odot\mathbf{A}$ is added to $\mathbf{X}_{i}\odot(1-\mathbf{A})$, hyper- or hypo-intensity is created in $\mathbf{X}$.

The mixing weight image $\bf{A}$ that determines the region of hyper-/hypo-intensity is designed to comprise both a tumor-like region and a buffer layer that represents the partial volume effect between normal tissue and regions with hyper-/hypo-intensity. To this end, we first generate a polyhedron with variable orientation and shape following [50]. The size of the polyhedron is sampled from the uniform distribution $U(20000~{}\mbox{mm}^{3},80000~{}\mbox{mm}^{3})$ [17, 50]. The center of the polyhedron is randomly placed within the brain, and the intersection of the inserted polyhedron and the brain mask is denoted by $\mathbf{M}$, which also represents the region with hyper-/hypo-intensity. To take the partial volume effect into consideration, 3D Gaussian filtering with a kernel size of 1 mm is applied to $\mathbf{M}$, and the filtered result $\bf{A}$ is used for the image generation in Eq. (1).

Brain tumors are generally segmented on T2-weighted or FLAIR images [2], where they appear as a cluster of hyper-intensity or also with a relatively dark component in the cluster. To generate images with tumor-like hyper-intensity, we apply Eq. (1) with $\lambda>1$ sampled from the uniform distribution $U(1.5,5)$. To further obtain images with a dark component, after applying Eq. (1), we use the following equations for image generation

Here, $\widetilde{T}(\mathbf{X}_{i})$ is a different transformation $\widetilde{T}(\mathbf{X}_{i})=\tilde{\lambda}\cdot\mathbf{Blur}(\mathbf{X}_{i})$, where $\tilde{\lambda}$ is sampled from $U(0.8,1.2)$ to create the dark region (darker than the hyper-intensity); $\widetilde{\mathbf{A}}$ is another weight image that represents the dark component and is obtained like $\mathbf{A}$¹¹1The polyhedron associated with $\widetilde{\mathbf{A}}$ is constrained to be smaller than that associated with $\mathbf{A}$. Its center is randomly placed within $\mathbf{M}$, and its intersection with $\mathbf{M}$ is smoothed to obtain $\widetilde{\mathbf{A}}$.; $\widetilde{\mathbf{X}}$ and $\widetilde{\mathbf{Y}}$ are the synthetic image and the corresponding mask indicating abnormality, respectively. In this way, when Eqs. (1) and (3) are sequentially applied, tumor-like hyper-intensity with a dark component is generated.

By repeating the image generation procedure in Eq. (1) or using the combination of Eqs. (1) and (3) with randomly drawn images, a number of synthetic images with approximately tumor-shaped abnormal intensity can be generated (see the examples in Fig. 1 associated with “Training Set”). The probability of applying Eq. (1) alone or both Eqs. (1) and (3) is set to 0.9 or 0.1, respectively. These synthetic images and the corresponding masks of abnormality are used to train the prompt model.

The abnormality generated with the procedures above can be too simplistic compared with real tumors. Thus, the prompt model obtained after training convergence can overfit the simplistic synthetic abnormality and generalize poorly to indicate real brain tumors. Ideally, to avoid overfitting, a validation set with annotated data is used to monitor the training process [6]. However, in the unsupervised setting there is no annotated data as the validation set.

To address this problem, we assume that during model training, the model first learns generic rules that apply to both simplistic synthetic abnormality and real tumors and then starts to overfit the detailed texture of synthetic abnormality [27]. Based on the assumption, we propose to use a validation task to prevent overfitting. This validation task also aims to identify abnormal intensity generated with hand-crafted rules, but the design of the rules is different from the prompt task for training, so that the two tasks require similar generic knowledge but the task details are different. In this way, after the generic knowledge about identifying abnormal intensity is sufficiently learned and the prompt model starts to overfit, the performance of the validation task will degrade, and based on this model selection can be performed.

A schematic of the benefit of the validation task is shown in Fig. 2. In the beginning epochs, the prompt model learns useful knowledge from the synthetic data, where both the training loss decreases and the accuracy of brain tumor segmentation on test data increases. Later, the model starts to overfit, where the training loss still decreases but the test accuracy starts to drop rapidly. Without the validation task, model selection occurs after the turning point of the test accuracy and is associated with overfitting, even if a validation set comprising the synthetic data is split from the training set for model selection; with the validation task, the turning point can be better identified for model selection.

The detailed design of the validation task is described as follows, which is similar to yet different from the prompt task. Specifically, we replace the transformation $T(\mathbf{X}_{i})$ and weight image $\mathbf{A}$ in Eq. (1) with

Compared with the prompt task, in the validation task no blurring is performed for image generation, and thus the validation task is similar to the prompt task but has different image textures near and in the tumors. Note that the set $\mathcal{X}$ is split into two subsets for the prompt task and validation task.

Since overfitting is more likely to occur when more training epochs are used and the scheduling of the optimization algorithm, such as Adam [51], is dependent on the maximum number $T$ of epochs, we considered different choices of $T$ for training the prompt model. Specifically, the set of candidate values of $T$ is $\{100,50,20,10,5\}$, and $T$ is also selected based on the performance of the validation task.

The OASIS dataset [53] is a publicly available dataset for neuroimaging study and analysis. There are no images with brain tumors in the dataset, which meets our need for prompt learning. We randomly selected 100 FLAIR images from the dataset and generated one image with abnormality from each of them (100 in total) for the prompt task; 100 different FLAIR images were also selected from the dataset, and one image with abnormality was generated from each of them (100 in total) for the validation task. The FLAIR images have been registered to the corresponding T1-weighted images and skull-stripped with BET [56]. All these images have the same voxel size of 1 mm isotropic.

The BraTS2021 dataset [54] comprises multimodal MR images with brain tumors acquired from multiple institutions under standard clinical conditions. There are 1251 cases of which both the images and tumor annotations are publicly available. The multimodal images of the same subject have been co-registered and skull-stripped. All these images have the same voxel size of 1 mm isotropic. In our experiment, the FLAIR images were selected for the segmentation of the whole tumor. All 1251 cases were used as the test set to evaluate the performance of the prompt model, as no annotation was required to train the prompt model.

As described in Section 3.2, our method allows the incorporation of unlabeled images with brain tumors into the training process. To evaluate the benefit of the incorporation, 100 images were randomly selected from the 1251 cases and included in training without using their annotations. The pseudo-labels predicted by the prompt model on the 100 images were pasted onto the 200 FLAIR images from the OASIS dataset. Note that the combination of the pseudo-label and the OASIS image was randomly selected. We ensured that each of the 200 FLAIR images was used once and each pseudo-label was used twice. This led to 200 synthetic images with pseudo-labels for model fine-tuning. In this case, the remaining 1151 cases were used as the test set.

BTFLAIR is an in-house dataset provided by Beijing Tiantan Hospital for segmenting whole brain tumors (ethical approval reference number: KY2022-078-04). The dataset includes 67 annotated FLAIR images acquired on multiple scanners, and they have been skull-stripped with BET [56]. The voxel size of these images ranges from $0.875\mbox{ mm}\times 0.875\mbox{ mm}\times 2\mbox{ mm}$ to $2\mbox{ mm}\times 2\mbox{ mm}\times 5\mbox{ mm}$. Like the BraTS2021 dataset, all images in BTFLAIR were used as the test set to evaluate the performance of the prompt model.

BTDWI is an in-house DWI dataset provided by Beijing Tiantan Hospital (ethical approval reference number: 2021-KY-0112). There are 47 DWIs acquired with a $b$-value of 1000 s/$\mbox{mm}^{2}$, and the skull was removed for each DWI with BET [56]. The voxel size of the DWIs is around $0.5\mbox{ mm}\times 0.5\mbox{ mm}\times 2\mbox{ mm}$. There are no ischemic stroke lesions on the DWIs, and the 47 images were combined with the three images without ischemic stroke lesions described next in Section 4.1.5 (50 in total) to train the prompt model that segments ischemic stroke lesions. 25 of the 50 DWIs without ischemic stroke lesions were used for the prompt task, and the other 25 DWIs were used for the validation task. We generated four images with abnormality from each DWI. Note that as the size of ischemic stroke lesions is different from that of brain tumors [57], during image generation the size of the polyhedron was adjusted to be sampled from a mixture of two uniform distributions $\frac{1}{2}U(200~{}\mbox{mm}^{3},300~{}\mbox{mm}^{3})+\frac{1}{2}U(5000~{}\mbox{mm}^{3},30000~{}\mbox{mm}^{3})$; also, as ischemic stroke lesions appear as hyper-intensity on DWIs, no dark component was used for image generation.

The ISLES2022 dataset [55] is publicly available and comprises 250 cases of ischemic stroke patients. We used the skull-stripped DWIs in the dataset, which were acquired on different scanners with a $b$-value of 1000 s/$\mbox{mm}^{2}$. The voxel size of these images ranges from $0.23\mbox{ mm}\times 0.23\mbox{ mm}\times 6\mbox{ mm}$ to $0.86\mbox{ mm}\times 0.86\mbox{ mm}\times 7\mbox{ mm}$. There are three cases that contain no ischemic stroke lesions. They were used for training the prompt model together with the BTDWI dataset as described in Section 4.1.4. The other 247 DWIs that contain ischemic stroke lesions were all used as the test set to evaluate the performance of the prompt model.

Again, to further investigate the benefit of incorporating unlabeled images with ischemic stroke lesions in training, we randomly selected 100 DWIs from the 247 DWIs with ischemic stroke lesions and added them into the training process as described in Section 3.2 without using their annotations. The pseudo-label predicted by the prompt model on each of the 100 DWIs was randomly pasted onto the 50 DWIs without ischemic stroke lesions from the BTDWI and ISLES2022 datasets. We ensured that each pseudo-label was used twice and each DWI without ischemic stroke lesions was used four times. This resulted in 200 synthetic images with ischemic stroke lesions for model fine-tuning. In this case, the remaining 147 DWIs were used as the test set.

First, for qualitative evaluation, representative segmentation examples selected from the test sets of BraTS2021 and BTFLAIR are shown in Figs. 4a and 4b, respectively, where the results of PL-BTS and the competing unsupervised methods are compared. The manual annotation is also displayed for reference. For both datasets our method produced tumor segmentation that is consistent with the annotation and better resembles the annotation than the results of the competing methods.

To quantitatively compare PL-BTS and each competing method, we computed the Dice coefficient, voxelwise precision, and voxelwise recall of the segmentation results on the test set for BraTS2021 and BTFLAIR. The means and standard deviations of these evaluation metrics are listed in Table 1 for the two datasets separately, where paired Student’s $t$-tests were also used to compare the results of PL-BTS and each competing method. For both datasets and all three metrics, the result of PL-BTS is highly significantly ($p<0.001$) better than those of the competing methods. Compared with the second best method CUAS, PL-BTS has improved the Dice coefficient by about 45.3% for BraTS2021 and 59.8% for BTFLAIR.

In addition, to show the necessity of the validation task in PL-BTS, the quantitative results of PL-BTS achieved without the validation task (referred to as PL-BTS w/o VT) are shown in Table 1. Compared with PL-BTS, the Dice coefficients of PL-BTS w/o VT drop by a large amount, which demonstrates the importance of using the validation task.

For qualitative evaluation, representative segmentation cases for the test set of ISLES2022 are shown in Fig. 5 together with the manual annotation, where the results of PL-BTS and the competing methods are compared. The results of PL-BTS in general agree with the annotation and better resemble the annotation than those of the competing methods.

For quantitative evaluation, the means and standard deviations of the Dice coefficient, voxelwise precision, and voxelwise recall of the segmentation results on the test set are shown in Table 2, where PL-BTS was compared with all competing methods with paired Student’s $t$-tests. The performance of PL-BTS is highly significantly ($p<0.001$) better than those of the competing methods in all cases. Compared with the second best method CUAS, PL-BTS has improved the Dice coefficient by about 141.4%.