Breast cancer detection and classification in mammogram using a three-stage deep learning framework based on PAA algorithm

Abstract

In recent years, deep learning has been used to develop an automatic breast cancer detection and classification tool to assist doctors. In this paper, we proposed a three-stage deep learning framework based on an anchor-free object detection algorithm, named the Probabilistic Anchor Assignment (PAA) to improve diagnosis performance by automatically detecting breast lesions (i.e., mass and calcification) and further classifying mammograms into benign or malignant. Firstly, a single-stage PAA-based detector roundly finds suspicious breast lesions in mammogram. Secondly, we designed a two-branch ROI detector to further classify and regress these lesions that aim to reduce the number of false positives. Besides, in this stage, we introduced a threshold-adaptive post-processing algorithm with dense breast information. Finally, the benign or malignant lesions would be classified by an ROI classifier which combines local-ROI features and global-image features. In addition, considering the strong correlation between the task of detection head of PAA and the task of whole mammogram classification, we added an image classifier that utilizes the same global-image features to perform image classification. The image classifier and the ROI classifier jointly guide to enhance the feature extraction ability and further improve the performance of classification. We integrated three public datasets of mammograms (CBIS-DDSM, INbreast, MIAS) to train and test our model and compared our framework with recent state-of-the-art methods. The results show that our proposed method can improve the diagnostic efficiency of radiologists by automatically detecting and classifying breast lesions and classifying benign and malignant mammograms.