Dual-Stage Deep Learning Framework for Prostate Cancer Grading Using Swin U-Net and Attention-Based CNNs

Abstract

Accurate grading of prostatic adenocarcinoma is essential in treatment planning. However, Gleason grading is time-consuming and clinically undependable. We presented a hybrid deep learning framework which comprises Swin U-Net for transformer-based segmentation network and attention-based CNNs for ISUP grade classification task. We incorporated Grad-CAM to aid in model interpretability and to visualize decision crucial areas. Quantitative evaluations on the PANDA, ISUP Grade-wise and transverse datasets achieve 100% accuracy on the smaller balanced Transverse dataset, 90.2 ± 0.7% performance in terms of ISUP with only 3.5M parameters, and a vicious Dice score equal to 0.99 ± 0.005 for segmentation. Notably, this cross-dataset generalization has not deteriorated below 92.3 ± 1.4% in any TIO experiment with no form of retraining applied to the transferred models. Inference time is less than 20 ms, deployment on the edge and mobile. The proposed model has achieved state-of-the-art performance for interpretability, accuracy, and computational complexity. The broadcast-then-categorize platform has been validated in ablation and optimization experiments, which demonstrate the potential for real-time diagnosis of prostate cancer.