[论文解读] Improved Unet model for brain tumor image segmentation based on ASPP-coordinate attention mechanism

Zixuan Wang, Yanlin Chen+2 more|arXiv (Cornell University)|Sep 13, 2024

Brain Tumor Detection and ClassificationNeuroscience被引用 3

一句话总结

本文提出了一种改进的U-Net模型，通过将空洞空间金字塔池化（ASPP）模块与坐标注意力机制相结合，显著提升了特征表示能力和边界定位精度。该模型在脑肿瘤分割任务中达到0.76的峰值mIoU，优于基线U-Net，在医学影像数据集上展现出更优的分割性能。

ABSTRACT

In this paper, we propose an improved Unet model for brain tumor image segmentation, which combines coordinate attention mechanism and ASPP module to improve the segmentation effect. After the data set is divided, we do the necessary preprocessing to the image and use the improved model to experiment. First, we trained and validated the traditional Unet model. By analyzing the loss curve of the training set and the validation set, we can see that the loss value continues to decline at the first epoch and becomes stable at the eighth epoch. This process shows that the model constantly optimizes its parameters to improve performance. At the same time, the change in the miou (mean Intersection over Union) index shows that the miou value exceeded 0.6 at the 15th epoch, remained above 0.6 thereafter, and reached above 0.7 at the 46th epoch. These results indicate that the basic Unet model is effective in brain tumor image segmentation. Next, we introduce an improved Unet algorithm based on coordinate attention mechanism and ASPP module for experiments. By observing the loss change curves of the training set and the verification set, it is found that the loss value reaches the lowest point at the sixth epoch and then remains relatively stable. At the same time, the miou indicator has stabilized above 0.7 since the 20th epoch and has reached a maximum of 0.76. These results show that the new mechanism introduced significantly improves the segmentation ability of the model. Finally, we apply the trained traditional Unet model and the improved Unet model based on the coordinate attention mechanism and ASPP module to the test set for brain tumor image segmentation prediction. Compared to the traditional Unet, the enhanced model offers superior segmentation and edge accuracy, providing a more reliable method for medical image analysis with the coordinate attention mechanism and ASPP module.

研究动机与目标

• 通过深度学习提升MRI图像中脑肿瘤分割的准确性。
• 解决标准U-Net在捕捉多尺度上下文信息和建模长距离依赖关系方面的局限性。
• 通过注意力机制与多速率卷积增强特征表示与边缘检测能力。
• 在公开脑肿瘤数据集上验证所提出架构的有效性。
• 为临床级医学图像分割提供一种稳健、可靠的解决方案。

提出的方法

• 所提出的模型集成空洞空间金字塔池化（ASPP）模块，用于从编码器特征中提取多尺度上下文特征。
• 引入坐标注意力机制，通过同时关注通道与空间坐标，显式建模空间依赖关系。
• 将ASPP与坐标注意力模块融合于U-Net的编码器-解码器架构中，以增强特征学习能力。
• 在预处理阶段应用数据增强与归一化操作，以提升模型泛化能力。
• 采用二元交叉熵损失函数，配合Adam优化器，并基于验证损失实施早停策略进行训练。
• 在测试集上进行推理，通过mIoU与视觉质量评估分割性能。

实验结果

研究问题

• RQ1 与标准U-Net相比，ASPP与坐标注意力的融合是否能提升脑肿瘤分割的mIoU？
• RQ2 所提出的注意力机制是否能增强特征表示能力并提升肿瘤区域的边界定位精度？
• RQ3 与基线U-Net相比，改进模型的训练收敛性与稳定性如何？
• RQ4 该模型在多大程度上减少了脑肿瘤分割中的假阳性和假阴性？
• RQ5 该增强模型在未见测试数据上是否具备良好的泛化能力，并实现高精度的边缘分割？

主要发现

• 改进后的U-Net模型达到0.76的峰值mIoU，显著优于基线U-Net。
• 改进模型的mIoU自第20个周期起稳定高于0.7，表明具有良好的收敛性。
• 改进模型的损失曲线在第6个周期达到最小值并保持稳定，表明训练速度更快且更稳定。
• 基线U-Net仅在第46个周期后mIoU才超过0.65，表明收敛速度更慢。
• 视觉检查确认，改进模型生成的肿瘤边界更清晰，分割噪声更少。
• 在测试集上，与标准U-Net相比，所提模型展现出更优的边缘精度与整体分割质量。