A lightweight Convolutional Neural Network based on U shape structure and Attention Mechanism for Anterior Mediastinum Segmentation

TL;DR

This paper presents a lightweight U-shaped CNN with attention mechanisms for accurate segmentation of anterior mediastinum lesions, improving diagnostic support in low-prevalence cases.

Contribution

Introduces a novel lightweight U-shaped network with Wide-MHSA and DDWPP modules for enhanced anterior mediastinum segmentation.

Findings

Achieved an average DSC of 87.83% on 2775 cases

Outperformed existing segmentation networks like TransUnet and Attention Unet

Demonstrated high sensitivity of 89.60%

Abstract

To automatically detect Anterior Mediastinum Lesions (AMLs) in the Anterior Mediastinum (AM), the primary requirement will be an automatic segmentation model specifically designed for the AM. The prevalence of AML is extremely low, making it challenging to conduct screening research similar to lung cancer screening. Retrospectively reviewing chest CT scans over a specific period to investigate the prevalence of AML requires substantial time. Therefore, developing an Artificial Intelligence (AI) model to find location of AM helps radiologist to enhance their ability to manage workloads and improve diagnostic accuracy for AMLs. In this paper, we introduce a U-shaped structure network to segment AM. Two attention mechanisms were used for maintaining long-range dependencies and localization. In order to have the potential of Multi-Head Self-Attention (MHSA) and a lightweight network, we designed a parallel MHSA named Wide-MHSA (W-MHSA). Maintaining long-range dependencies is crucial for segmentation when we upsample feature maps. Therefore, we designed a Dilated Depth-Wise Parallel Path connection (DDWPP) for this purpose. In order to design a lightweight architecture, we introduced an expanding convolution block and combine it with the proposed W-MHSA for feature extraction in the encoder part of the proposed U-shaped network. The proposed network was trained on 2775 AM cases, which obtained an average Dice Similarity Coefficient (DSC) of 87.83%, mean Intersection over Union (IoU) of 79.16%, and Sensitivity of 89.60%. Our proposed architecture exhibited superior segmentation performance compared to the most advanced segmentation networks, such as Trans Unet, Attention Unet, Res Unet, and Res Unet++.