Unet-boosted Classifier – Multi-Task Architecture for Small Datasets Applied to Brain MRI Classification

Abstract

The problem of training deep neural networks on small samples is especially relevant for medical issues. The paper examines the impact of pixel-wise marking of significant objects in the image, over the true class label, on the quality of the classification. To achieve better classification results on small samples, we propose a multitasking architecture – Unet-boosted classifier (UBC), that is trained simultaneously to solve classification and semantic segmentation problems. As the exploratory dataset, MRI images of patients with benign glioma and glioblastoma taken from the BRaTS 2019 data set are used. One horizontal slice of the MRI image containing a glioma is considered as the input (a total of 380 frames in the training set), and the probability of glioblastoma – as the output. Resnet34 was chosen as the baseline, trained without augmentations with a loss function based on cross-entropy. As an alternative solution, UBC-resnet34 is used – the same resnet34, boosted by a decoder built on the U-Net principle and predicting the pixels with glioma. The smoothed Sorensen-Dice coefficient (DiceLoss) is used as a decoder loss function. Results on the test sample: accuracy for the baseline reached 0.71, for the proposed model – 0.81, and the Dice score – 0.77. Thus, a deep model can be well trained even on a small data set, using the proposed architecture, provided that marking of the affected tissues in the form of a semantic mask is provided.