Multi-stage Fine-tuning Approach for AI-based Chest X-ray Abnormality Detection

Abstract

Correct detection and localization of thoracic diseases on chest X-ray images supports or even secures an early diagnosis and treatment-oriented planning. This paper proposes a version of the YOLOv5 deep learning model augmented with more sophisticated components, namely Spatial Pyramid Pooling (SPP), Path Aggregation Network (PANet), and Convolutional Block Attention Module (CBAM) to enhance overall stability and generalization for clinical use. The key components of the proposed model are first trained on the VinBigData Chest X-ray Abnormalities Detection dataset to improve feature extraction and adaptively stretch multiple image resolutions across different strides per pixel levels while focusing more attention on the region. Next, we make a multi-stage fine-tuning approach for real-world clinical data, which usually shifts domains in practical settings. Finally, the model is forced to be more resistant and less overfitting by performing real-world data augmentation instead of focusing on clinical variability. We further qualitatively assess the performance of our model on both the VinBigData test set and CheXDet dataset with only publicly available bounding box annotations on matching classes between the two datasets. Moreover, the model was integrated into a web application that could easily be employed in clinical environments for real-time chest X-ray analysis and may assist with more accurate diagnosis at an early stage.