Extending nn-UNet for brain tumor segmentation

TL;DR

This paper enhances the nn-UNet architecture for brain tumor segmentation by experimenting with modifications like larger networks, group normalization, and axial attention, leading to a winning solution in the 2021 BraTS challenge.

Contribution

The paper introduces specific modifications to nn-UNet, including larger networks and attention mechanisms, demonstrating improved performance in brain tumor segmentation.

Findings

Achieved first place in the BraTS 2021 challenge.

Minor improvements in quantitative metrics over baseline.

Models and code are publicly available.

Abstract

Brain tumor segmentation is essential for the diagnosis and prognosis of patients with gliomas. The brain tumor segmentation challenge has continued to provide a great source of data to develop automatic algorithms to perform the task. This paper describes our contribution to the 2021 competition. We developed our methods based on nn-UNet, the winning entry of last year competition. We experimented with several modifications, including using a larger network, replacing batch normalization with group normalization, and utilizing axial attention in the decoder. Internal 5-fold cross validation as well as online evaluation from the organizers showed the effectiveness of our approach, with minor improvement in quantitative metrics when compared to the baseline. The proposed models won first place in the final ranking on unseen test data. The codes, pretrained weights, and docker image for the winning submission are publicly available at https://github.com/rixez/Brats21_KAIST_MRI_Lab