Addressing class imbalance in deep learning for small lesion detection on medical images

Abstract

Deep learning methods utilizing Convolutional Neural Networks (CNNs) have led to dramatic advances in automated understanding of medical images. However, in many medical image classification tasks, lesions occupy only a few pixels of the image. This results in a significant class imbalance between lesion and background. From recent literature, it is known that class imbalance may negatively affect the performance of CNN classification. However, very few research exists in the context of lesion detection. In this work, we propose a two-stage deep learning framework able to deal with the high class imbalance encountered during training of small lesion detectors. First, we train a deep cascade (DC) of long sequences of decision trees with an algorithm designed to handle unbalanced data that also drastically reduces the number of background samples reaching the final stage. The remaining samples are fed to a CNN, whose training benefits from both rebalance and hard mining done by the DC. We evaluated DC-CNN on two severely unbalanced classification problems: microcalcification detection and microaneurysm detection. In both cases, DC-CNN outperformed the CNNs trained with commonly used methods for addressing class imbalance such as oversampling, undersampling, hard mining, cost sensitive learning, and one-class classification. The DC-CNN was also ∼10x faster than CNN at test time.

Keywords: Class imbalance; Deep learning; Lesion detection; Microaneurysms; Microcalcifications.