# YM-WML: A new Yolo-based segmentation Model with Weighted Multi-class Loss for medical imaging

arXiv: https://arxiv.org/abs/2506.22955 · HTML: https://mlsanity.com/paper/2506.22955/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 14 of its 24 references that are in the map.
- First posted on arXiv 28 Jun 2025.

## Builds on (most cited first, 14 of 14)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Fully Convolutional Networks for Semantic Segmentation](https://mlsanity.com/paper/1411.4038/) (2014)
- [You Only Look Once: Unified, Real-Time Object Detection](https://mlsanity.com/paper/1506.02640/) (2015)
- [MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications](https://mlsanity.com/paper/1704.04861/) (2017)
- [Attention U-Net: Learning Where to Look for the Pancreas](https://mlsanity.com/paper/1804.03999/) (2018)
- [MaxViT: Multi-Axis Vision Transformer](https://mlsanity.com/paper/2204.01697/) (2022)
- [A survey of loss functions for semantic segmentation](https://mlsanity.com/paper/2006.14822/) (2020)
- [Boundary loss for highly unbalanced segmentation](https://mlsanity.com/paper/1812.07032/) (2018)
- [LeViT-UNet: Make Faster Encoders with Transformer for Medical Image Segmentation](https://mlsanity.com/paper/2107.08623/) (2021)
- [A review: Deep learning for medical image segmentation using multi-modality fusion](https://mlsanity.com/paper/2004.10664/) (2020)
- [Understanding Deep Learning Techniques for Image Segmentation](https://mlsanity.com/paper/1907.06119/) (2019)

## Abstract (arXiv)

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.