Region-related focal loss for 3D brain tumor MRI segmentation

In the brain tumor magnetic resonance image (MRI) segmentation, although the 3D convolution networks (CNNs) has achieved state-of-the-art results, the class and hard-voxel imbalances in the 3D images have not been well addressed. Voxel independent losses are dependent on the setting of class weights for the class imbalance issue, and are hard to assign each class equally. Region-related losses cannot correctly focus on hard voxels dynamically and not be robust to misclassification of small structures. Meanwhile, repeatedly training on the additional hard samples augmented by existing methods would bring more class imbalance, overfitting and incorrect knowledge learning to the model.

A novel region-related loss with balanced dynamic weighting while alleviating the sensitivity to small structures is necessary. In addition, we need to increase the diversity of hard samples in the training to improve the performance of model.

The proposed Region-related Focal Loss (RFL) reshapes standard Dice Loss (DL) by up-weighting the loss assigned to hard-classified voxels. Compared to DL, RFL adaptively modulate its gradient with an invariant focalized point that voxels with lower-confidence than it would achieve a larger gradient, and higher-confidence voxels would get a smaller gradient. Meanwhile, RFL can adjust the parameters to set where and how much the network is focused. In addition, an Intra-classly Transformed Augmentation network (ITA-NET) is proposed to increase the diversity of hard samples, in which the 3D registration network and intra-class transfer layer are used to transform the shape and intensity respectively. A selective hard sample mining(SHSM) strategy is used to train the ITA-NET for avoiding excessive class imbalance. Source code (in Tensorflow) is available at: https://github.com/lb-whu/RFL_ITA.

The experiments are carried out on public data set: Brain Tumor Segmentation Challenge 2020 (BratS2020). Experiments with BraTS2020 online validation set show that proposed methods achieve an average Dice scores of 0.905, 0.821, and 0.806 for whole tumor (WT), tumor core (TC) and enhancing tumor (ET), respectively. Compared with DL (baseline), the proposed RFL significantly improves the Dice scores by an average of 1%, and for the small region ET it can even increase by 3%. And the proposed method combined with ITA-NET improves the Dice scores of ET and TC by 5% and 3% respectively.

The proposed RFL can converge with a invariant focalized point in the training of segmentation network, thus effectively alleviating the hard-voxel imbalance in brain tumor MRI segmentation. The negative region term of RFL can effectively reduce the sensitivity of the segmentation model to the misclassification of small structures. The proposed ITA-NET can increase the diversity of hard samples by transforming their shape and transfer their intra-class intensity, thereby effectively improving the robustness of the segmentation network to hard samples.