YM-WML: A new Yolo-based segmentation Model with Weighted Multi-class Loss for medical imaging
This is an incremental improvement for medical imaging researchers, specifically enhancing cardiac segmentation accuracy.
The paper tackles cardiac image segmentation by proposing YM-WML, a YOLO-based model with a novel Weighted Multi-class Exponential loss function to address class imbalance, achieving a Dice Similarity Coefficient of 91.02 on the ACDC dataset and outperforming state-of-the-art methods.
Medical image segmentation poses significant challenges due to class imbalance and the complex structure of medical images. To address these challenges, this study proposes YM-WML, a novel model for cardiac image segmentation. The model integrates a robust backbone for effective feature extraction, a YOLOv11 neck for multi-scale feature aggregation, and an attention-based segmentation head for precise and accurate segmentation. To address class imbalance, we introduce the Weighted Multi-class Exponential (WME) loss function. On the ACDC dataset, YM-WML achieves a Dice Similarity Coefficient of 91.02, outperforming state-of-the-art methods. The model demonstrates stable training, accurate segmentation, and strong generalization, setting a new benchmark in cardiac segmentation tasks.