One-Shot Traumatic Brain Segmentation with Adversarial Training and Uncertainty Rectification

Xiangyu Zhao1, Zhenrong Shen1, Dongdong Chen1, Sheng Wang1,3, Zixu Zhuang1,3, Qian Wang2, Lichi Zhang1,
1School of Biomedical Engineering, Shanghai Jiao Tong University, Shanghai, China
 2School of Biomedical Engineering, ShanghaiTech University, Shanghai, China
 3Shanghai United Imaging Intelligence Co., Ltd., Shanghai, China
MICCAI 2023
Paper Supplementary Poster Code

Abstract

Brain segmentation of patients with severe traumatic brain injuries (sTBI) is essential for clinical treatment, but fully-supervised segmentation is limited by the lack of annotated data. One-shot segmentation based on learned transformations (OSSLT) has emerged as a powerful tool to overcome the limitations of insufficient training samples, which involves learning spatial and appearance transformations to perform data augmentation, and learning segmentation with augmented images. However, current practices face challenges in the limited diversity of augmented samples and the potential label error introduced by learned transformations. In this paper, we propose a novel one-shot traumatic brain segmentation method that surpasses these limitations by adversarial training and uncertainty rectification. The proposed method challenges the segmentation by adversarial disturbance of augmented samples to improve both the diversity of augmented data and the robustness of segmentation. Furthermore, potential label error introduced by learned transformations is rectified according to the uncertainty in segmentation. We validate the proposed method by the one-shot segmentation of consciousness-related brain regions in traumatic brain MR scans. Experimental results demonstrate that our proposed method has surpassed state-of-the-art alternatives.

Overview

Automatic brain ROI segmentation for magnetic resonance images (MRI) of severe traumatic brain injuries (sTBI) patients is crucial in brain damage assessment and brain network analysis, since manual labeling is time-consuming and labor-intensive. However, conventional brain segmentation pipelines, such as FSL and FreeSurfer, suffer significant performance deteriorations due to skull deformation and lesion erosions in traumatic brains. Although automatic segmentation based on deep learning has shown promises in accurate segmentation, these methods are still constrained by the scarcity of annotated sTBI scans. Thus, researches on traumatic brain segmentation under insufficient annotations needs further exploration.

Traumatic Brain MR

Here is the illustration of the MR image and its consciousness-related brain regions of a sTBI patient. Compared with normal brain scans, brain segmentation of sTBI brain scans are affected by lesions and deformed skulls, which leads to deteriorated performance. In addition, manual labeling of these brain regions in sTBI MR scans could be laborious, and thus available annotations are extremly scarce.

One-Shot Segmentation

One-shot segmentation based on learned transformations (OSSLT) has been applied in one-shot brain segmentation. These methods typically include three basic steps: 1) Learning registration; 2) Atlas augmentation; 3) Learning segmentation. However, previous OSSLT methods are faced with two challenges. First, the diversity of atlas augmentation is limited by the amount of available images. Second, the label authenticity of generated images is affected by the presence of brain trauma during atlas augmentation. Thus, we propose to introduce adversarial training for atlas augmentation and uncertainty rectification for segmentation to address these issues.

Adversarial Training

We propose adversarial training for atlas augmentation to improve the diversity of augmented atlas images. The proposed adversarial training strategy consists of three steps. First, the adversarial network takes original generated image as input, and executes adversarial sampling to spatial and appearance transformations. After that, the adversarial transformations are applied to the atlas image to create an adversarial augmented image, which is fed to the segmentation network. Then, the adversarial network and the segmentation network is trained in an adversarial manner. To be specific, the adversarial network is trained to maximize the segmentation difference before and after applying adversarial transformations, while the segmentation network is trained to minimize the segmentation difference. Thus, proposed adversarial training is able to improve both the diversity of augmented atlas images and the robustness of segmentation learning.

Uncertainty Rectified Segmentation

The presence of brain trauma could affect the label authenticity of the generated image. Thus, we propose uncertainty rectification to alleviate this problem. Specifically, we use the warped atlas image as the spatial reference. The segmentation difference of the generated image and the warped atlas image may indicate the locations of potential label errors, which are then rectified by spatial-weighted segmentation loss. The weight of different locations is determined by the KL-divergence of two segmentation maps.

Segmentation Visualization

Left to Right: LT-Net, DeepAtlas, Brainstorm, Ours, and ground truth.

Quantitative Results

BibTeX

@InProceedings{Zhao_2023_MICCAI,
          author    = {Zhao, Xiangyu and Shen, Zhenrong and Chen, Dongdong and Wang, Sheng and Zhuang, Zixu and Wang, Qian and Zhang, Lichi},
          title     = {One-Shot Traumatic Brain Segmentation with Adversarial Training and Uncertainty Rectification},
          booktitle = {Proceedings of the International Conference on Medical Image Computing and Computer Assisted Intervention},
          month     = {October},
          year      = {2023},
        }