Seg-ViT-Opt: A Segmentation Enhanced Vision Transformer Framework with Optimized Training for Cervical Cell Classification
摘要
Cervical cancer ranks as the fourth most prevalent malignancy among females globally and is a major source of cancer-related mortality; it is especially so in low-and-middle income countries. Cytology based screening programs have been successful, but they are time consuming, subject to observer variability, and thus fail to enable early detection and timely treatment of cervical cancer. As such, deep learning has proven to be a powerful tool that can be used to solve problems associated with cytology diagnostics. In particular, Convolutional Neural Network (CNN) models have achieved great success in the automatic cytology image analysis process. However, CNNs are limited due to their local receptive fields; therefore, CNNs cannot effectively identify the global contextual relationship between different parts of a cell. More recent developments in medical imaging have included the introduction of Vision Transformers (ViTs), where ViTs offer improved representation capabilities compared to CNNs by utilizing self-attention mechanisms to extract contextual information from each pixel within the image. The aim of this study was to discover new Segmentation enhanced Vision Transformer Architecture, named Seg-ViT-Opt, for optimizing the classification of Cervical cell. The Seg-ViT-Opt framework uses a U-Net-derived nucleus segmentation unit to generate binary segmentation masks which are then used to construct segmentation-guided representations in both masked RGB and four-channel representations using cytology imagery. Following generating these segmentation guided representations, the classification process is accomplished using either convolutional neural network or Vision Transformer backbone. The classification architecture is further optimized through the use of six different optimizers (i.e., AdamW, RAdam, Lion, Ranger, SGD, and AdamW with Sharpness Aware Minimization). The proposed Seg-ViT-Opt model demonstrated superior classification performance on the SIPaKMeD dataset with a best state-of-the-art accuracy of 99.01% and a macro F1 score of 0.990 when utilizing the AdamW optimizer. Additionally, the Seg-ViT-Opt model demonstrated perfect classification for the Koilocytotic and Superficial Intermediate classes. The experimental results presented demonstrate that segmentation-guided ViTs represent a reliable, interpretable, and highly performing framework for the automation of cervical cancer diagnosis; furthermore, the results indicate that there is strong potential for the integration of segmentation-guided ViTs into clinical settings. This research differs from previous cytology-based studies in that it incorporates both segmentation-guided ViT input encoding and a broad benchmarking strategy for optimizers to provide a single, integrated and generalized framework for transformer-based cytology classification.