超越SOTA ! YOLOv8-ResCBAM 集成注意力机制以提高检测性能 !

科技   2024-10-17 09:00   上海  

点击下方卡片,关注「集智书童」公众号

点击加入👉「集智书童」交流群




想要了解更多:

前沿AI视觉感知全栈知识👉「分类、检测、分割、关键点、车道线检测、3D视觉(分割、检测)、多模态、目标跟踪、NerF

行业技术方案👉AI安防、AI医疗、AI自动驾驶
AI模型部署落地实战👉「CUDA、TensorRT、NCNN、OpenVINO、MNN、ONNXRuntime以及地平线框架」

欢迎扫描上方二维码,加入「集智书童-知识星球」,日常分享论文、学习笔记、问题解决方案、部署方案以及全栈式答疑,期待交流!

免责声明
凡本公众号注明“来源:XXX(非集智书童)”的作品,均转载自其它媒体,版权归原作者所有,如有侵权请联系我们删除,谢谢。

手腕创伤甚至骨折在日常生活中频繁发生,特别是在儿童中,他们占到了骨折病例的很大比例。在进行手术之前,医生通常会让病人先进行X光成像,并根据X光图像的分析来准备手术。

随着神经网络的发展,YOLO系列模型在计算机辅助诊断中的骨折检测中得到了广泛应用,其中YOLOv8模型已经取得了令人满意的结果。将注意力模块应用到神经网络中是一种有效的方法来提高模型性能。

本文提出了一种将具有resblock(ResCBAM)的卷积块注意力模块集成到原始YOLOv8网络架构中的YOLOv8-ResCBAM模型。

在GRAZPEDWRI-DX数据集上的实验结果表明,在交点与 union 阈值为0.5时的平均精确度(mAP 50)计算中,所提出的模型的平均精确度从原始YOLOv8模型的63.6%增加到65.8%,实现了最先进的效果。

代码https://github.com/RuiyangJu/Fracture_Detection_Improved_YOLOv8。

1 Introduction

手腕骨折是其中最常见的一种骨折,尤其是对于老年人以及儿童[16, 40],主要发生在距离关节末端2厘米处的前臂桡骨近端。如果不能及时进行治疗,可能导致手腕关节畸形、关节活动受限以及患者关节疼痛[2]。对于儿童来说,误诊可能导致终身的不便[26]。

在儿童手腕骨折的情况下,医生们通常会询问导致骨折的原因,并让患者进行骨折检查。

骨折检查主要通过三种医学影像设备进行:磁共振成像(MRI)、计算机断层扫描(CT)和X光。由于其成本效益,X光成为大多数患者的首选 [49]。在提供高级医疗护理的医院中,放射科医生需要遵循健康水平7(HL7)和医学影像与通信(DICOM)国际标准,对X光图像进行归档和传输 [5]。然而,发展中国家地区放射科医生的稀缺性给患者及时护理的提供带来了挑战 [6, 43, 44]。[12, 37]中的研究表明,在紧急情况下,医学影像分析的错误率令人担忧地达到了26% [4]。

计算机辅助诊断(CAD)为专家(放射科医生、外科医生等)在某些决策上提供了帮助。随着深度学习和医学影像处理技术的不断发展,越来越多的研究行人尝试使用神经网络进行CAD,包括骨折检测。

YOLO [41],作为目标检测任务中最重要的网络模型之一 [24],在骨折检测 [17]中表现出令人满意的模型性能。GRAZPEDWRI-DX [38] 是一个公开的骨折检测数据集,其中包含 20,327 张儿童手腕创伤 X 光图像。随着 YOLOv8 [21] 的引入,它在该数据集上用于计算机辅助设计(CAD) [23]。

由于注意力模块能够准确地关注输入图像的重要信息,它们被广泛应用于神经网络架构中。目前,主要有两种类型的注意力模块:空间注意力(spatial attention)和通道注意力(channel attention),分别用于捕捉像素 Level 的配对关系和通道依赖性。

研究表明,将注意力模块集成到卷积块中可以显著提高模型性能[7, 13, 30]。因此,本文提出YOLOv8-ResCBAM,将卷积块注意力模块(CBAM) [50] 与ResCBAM集成到原始YOLOv8网络架构中,在GRAZPEDWRI-DX数据集[38]上获得最先进(SOTA)的结果。

本文的主要贡献如下:

该方法提出了YOLOv8-ResCBAM用于儿童手腕骨折检测,通过将ResCBAM集成到原始YOLOv8网络架构中。

该模型在GRAZPEDWRI-DX数据集上显著优于YOLOv8模型,实现了最先进的表现。

它采用YOLOv8-ResCBAM作为计算机辅助诊断工具,以帮助外科医生分析手部X光图像。

2 Related Work

Fracture Detection

骨折检测是医学影像分析和处理的热门话题,研究行人通常在医学图像上使用神经网络,包括YOLO系列模型。在Nagy等人[38]发布GRAZPEDWRI-DX数据集之前,与儿科手腕X光图像相关的公开可用的儿科数据集较少。此外,成人病例的骨骼放射科资料也有限。

Hrzic等人[17]利用YOLOv4模型在GRAZPEDWRI-DX数据集上检测骨折,表明YOLO系列模型可以基于X光图像提高诊断儿童手腕损伤的准确性。

Ju等人[23]创建了"Fracture Detection Using YOLOv8 App",以帮助外科医生在X光图像上解释骨折,从而减少误诊风险,增加骨折手术所需的信息。

Attention Module

挤压与激励网络(SENet)[18] 最初提出了一种通过将全局平均池化(GAP)应用于每个通道来高效学习通道注意力的机制。随后,使用全连接层和 sigmoid 函数生成通道权重,从而实现了良好的模型性能。

在 SENet 引入特征聚合和特征校正之后,一些研究[8, 15]尝试通过捕获更复杂的通道依赖关系来改进 SE 块。Woo 等人[50] 将通道注意力模块与空间注意力模块相结合,引入了 CBAM,以提高卷积神经网络(CNNs)的表现能力。

这种注意力模块可以应用于不同的神经网络架构,以提高模型性能。

3 Methodology

Baseline

YOLOv8网络架构包括四个关键组件:Backbone、Neck、Head和Loss Function。Backbone集成了Cross Stage Partial(CSP)[47],这可以降低计算负荷,同时增强卷积神经网络的学习能力。如图1所示,YOLOv8与YOLOv5的区别在于使用了C3模块[22],并采用了C2f模块,该模块将C3模块和YOLOv7的Extended ELAN(E-ELAN)[48]集成在一起。

具体来说,C3模块包括三个卷积模块和多个 Bottleneck ,而C2f模块由两个卷积模块 ConCat 多个 Bottleneck 组成。卷积模块的结构为

在 Neck 分,YOLOv5采用了特征金字塔网络(FPN)[32]结构进行自上而下的采样,确保下层特征图包含更丰富的特征信息。同时,路径聚合网络(PAN)[34]结构用于自下而上的采样,为顶层特征图提供更精确的位置信息。将这两种结构相结合,以保证在不同尺寸的图像上实现准确的预测。YOLOv8遵循FPN和PAN框架,在向上采样阶段删除了卷积操作,如图1所示。

与YOLOv5相比,YOLOv8采用了解耦头,将分类头和检测Head分离。具体来说,YOLOv8去除了目标性分支,只保留了分类和回归分支。此外,它放弃了基于 Anchor 点的方法[42],转而采用 Anchor-Free 点方法[11],其中目标的位置由其中心确定,预测涉及估算中心到边界的距离。

在YOLOv8中,分类分支所采用的损失函数为二进制交叉熵(BCE)损失,其表达式如下:


其中表示权重;表示标签值,表示模型生成的预测值。

对于回归分支,YOLOv8采用了Distribute Focal Loss (DFL) [31]和Complete Intersection over Union (CIoU) Loss [56]。DFL函数旨在强调对目标的概率值的扩展,其方程如下所示:


引入了影响因素的CIoU损失将 DIoU 损失 [55] 引入了一个考虑预测框和真实框的宽高比的因素,相应的方程如下:


在这里, 衡量了预测框和真实框之间的重叠程度; 是预测框和真实框中心点之间的欧几里得距离, 是包含预测框和真实框的最小外接框的对角线长度。此外, 表示了用于量化长宽比一致性的参数,其定义如下方程所示:


其中,表示边界框的权重;表示边界框的高度;表示真实值,表示预测值。

Proposed Method

近年来,注意力模块在目标检测领域取得了卓越的成果 [20, 28, 53]。通过集成注意力模块,网络能够识别输入图像中最重要信息以提取,并抑制无用的信息。

这篇论文将有效的注意力模块集成到YOLOv8的Neck部分,以增强对关键特征的捕捉并抑制干扰信息。如图1所示,作者在每个C2f模块后添加了一个ResCBAM,因为CBAM已经证明在各种模型中都能一致提高分类和检测性能。

CBAM 包含两种注意力机制:通道注意力(C-Attention)和空间注意力(S-Attention),如图2所示。首先,从输入特征图 开始,CBAM 生成一个1D的通道注意力映射 ,然后生成一个2D的空间注意力映射 ,如方程所示:


其中表示元素乘法;是通道细化特征,是最终细化特征。对于CBAM,等于,如下所示的方程:

从图2可以看出,对于ResCBAM,的元素求和,如下所示:

根据先前的研究[39, 52],CBAM采用全局平均池化(GAP)和全局最大池化(GMP)来聚合特征图的空间信息,从而生成两个不同的空间上下文描述符。随后,这两个描述符共享同一个多层感知机(MLP)的一个隐藏层。最后,将元素逐个求和得到的特征向量输入到 sigmoid 函数()。具体的通道注意力方程如下:


对于空间注意力,CBAM 分别对通道轴进行 GAP 和 GMP,然后将它们拼接在一起,接下来,使用一个 卷积层对这些特征进行卷积操作。卷积操作的输出作为 的输入。空间注意力使用以下方程计算:


4 Experiment

Dataset

GRAZPEDWRI-DX [38] 是一个公开的儿科手腕创伤X光图像数据集,包含20,327张X光图像。这些X光图像是由奥地利格拉兹大学医院儿科手术部多位儿科放射科医生在2008年和2018年间收集的,涉及6,091名患者,共计10,643个研究。

该数据集附带74,459个图像标签,共有67,771个标注物体。该数据集表现出类不平衡,其中“文本”类标签有23,722个,“骨折”类标签有18,090个,而“骨异常”和“软组织”类标签分别只有276个和464个。

Preprocessing and Data Augmentation

在没有预定义的训练、验证和测试集的情况下,作者进行了随机的划分,其中70%分配到训练集,20%分配到验证集,10%分配到测试集。具体而言,训练集包括14,204张图像(69.88%);验证集包括4,094张图像(20.14%);测试集包括2,029张图像(9.98%)。

由于训练集中的X射线图像亮度多样性有限,仅在这些图像上训练的模型在预测其他X射线图像时可能表现不佳。此外,该数据集存在类不平衡问题。为了解决这些问题,作者采用数据增强技术来扩充训练集。

具体而言,作者使用开源计算机视觉库(OpenCV)中的addWeighted函数来微调X射线图像的对比度和亮度。

Evaluation Metric

在本研究中,作者在实际诊断场景中评估了YOLOv8-ResCBAM和YOLOv8模型在参数(Params)、浮点运算(FLOPs)、F1评分、0.5(mAP 50)的交并集平均精度,以及0.5至0.95(mAP 50-95)的交并集平均精度,以及推理时间。

Experiment Setup

作者在数据集[38]上训练了YOLOv8模型和YOLOv8-ResCBAM模型。与Ultralytics[21]为YOLOv8训练推荐的前300个周期不同,实验结果[23]表明,最佳性能在60至70个周期内实现。因此,本工作将所有模型的训练周期设置为100个周期。

对于模型训练的超参数,作者根据[23]中的消融研究结果,选择SGD优化器而不是Adam优化器。遵循Ultralytics[21]的建议,本工作将优化器的权重衰减设置为,配合0.937的动量,以及初始学习率设置为。为了比较不同输入图像尺寸对模型性能的影响,本工作将输入图像尺寸分别设置为640和1024进行实验。

本文采用Python 3.9在支持PyTorch 1.13.1框架的模型上进行训练。作者建议读者使用高于Python 3.7和PyTorch 1.7的版本进行模型训练,具体所需环境可在作者的GitHub仓库中找到。所有实验均在一个单独的NVIDIA GeForce RTX 3090 GPU上运行,批处理大小设置为16以适应GPU内存限制。

Ablation Study

在消融研究中,作者研究了ResCBAM对模型性能的影响,其中输入图像的大小不同。具体来说,作者分别使用输入图像大小为640和1024的训练集训练作者的模型,并在具有相应图像尺寸的测试集上评估YOLOv8-ResCBAM模型。消融研究的结果在表1中展示。YOLOv8-ResCBAM模型实现的F1分数和mAP 50值超过了YOLOv8(基础)模型的水平。

具体而言,当输入图像大小为1024时,YOLOv8-ResCBAM-L模型的mAP 50从63.6%提高到65.8%,相较于YOLOv8-L模型提高了3.5%。此外,YOLOv8-ResCBAM模型的推理时间(8.7ms/图像)与YOLOv8模型(7.8ms/图像)相当。

Experimental Results

如图1所示,使用输入图像大小为1024的模型训练的性能超过了使用相同训练集和输入图像大小为640的模型。然而,值得注意的是,模型性能的提高伴随着推理时间的增加。例如,对于大型模型大小的YOLOv8-ResCBAM模型,当输入图像大小为1024时,mAP 50-95值达到42.2%,比输入图像大小为640时的40.1%提高了5.24%。然而,该模型的推理时间从4.1ms增加到8.7ms。

为了证明作者提出的模型在GRAZPEDWRI-DX数据集上实现了最先进的性能,作者在输入尺寸为1024且大型模型大小条件下,与不同模型进行了比较,结果如表2所示。在提出YOLOv8-ResCBAM之前,该数据集上训练的模型的最高mAP 50值为65.0%,最高mAP 50-95值为41.9%。

然而,作者的YOLOv8-ResCBAM模型的mAP 50和mAP 50-95值分别为65.8%和42.2%,优于当前最先进的水平。

这篇论文评估了CBAM对YOLOv8模型在实际案例诊断场景中预测骨折准确性的影响。共选定了8张X光图像进行评估。图3展示了YOLOv8-ResCBAM和YOLOv8模型的预测结果,表明作者的骨折检测能力优于YOLOv8模型,尤其是在单处骨折的情况下。例如,在图3中的第四张X光图像中,YOLOv8预测骨折概率为30%,而所提出的模型预测概率为50%,正确地识别出地层图像中的骨折位置。因此,YOLOv8-ResCBAM模型可以作为CAD工具,帮助放射科医生和外科医生进行诊断。

图4说明了YOLOv8-ResCBAM模型对每个类别的预测精度-召回曲线。分析表明,YOLOv8-ResCBAM模型和YOLOv8模型都可以正确检测到骨折、金属和文本,平均准确率达到90%以上。

然而,它们检测“骨异常”和“软组织”的能力显著较弱,这极大地影响了模型在mAP 50值上的表现。具体而言,YOLOv8-ResCBAM模型在预测这两个类别时的准确率分别达到15.0%和31.5%,而YOLOv8仅为8.4%和25.6%。作者认为这是由于在所使用的数据集中,这两个类别的目标数量较少。

正如GRAZPEDWRI-DX [38]中描述的那样,骨异常和软组织的数量分别占全部目标数量的0.41%和0.68%。因此,通过架构优化提高模型性能受到这些数据限制的约束。为了增强模型性能,迫切需要采用额外的数据进行补充。

5 Conclusion and Future Work

随着YOLOv8的引入,研究行人开始将其应用于身体各部位的骨折检测。尽管YOLOv8模型在GRAZPEDWRI-DX数据集上的性能可圈可点,但尚未达到SOTA水平。为了克服这一局限,作者将ResCBAM集成到YOLOv8架构中,以提高模型性能。值得注意的是,当输入图像大小为1024时,YOLOv8-ResCBAM模型的mAP 50取得了65.8%的优越性能,超过了SOTA基准。此外,作者提出的YOLOv8-ResCBAM模型可以作为CAD工具,帮助外科医生分析X光图像,从而减少骨折检测中的误判可能性。

在未来的工作中,作者将尝试与医院或医学院合作,收集更多不寻常的儿童手腕X光图像(包括骨折或其他疾病)的标签(如“骨异常”和“软组织”),以提高模型性能。

作者希望作者的提出的YOLOv8-ResCBAM模型能够在未来帮助专家(放射科医生、外科医生等)做出更好的决策。

具体来说,在作者的先前的研究中[23],作者引入了“使用YOLOv8的骨折检测”,作者计划将提出的YOLOv8-ResCBAM模型集成到这个应用中。

扫码加入👉「集智书童」交流群

(备注:方向+学校/公司+昵称



点击下方“阅读原文”,
了解更多AI学习路上的「武功秘籍」


集智书童
书童带你领略视觉前沿之美,精选科研前沿、工业实用的知识供你我进步与学习!
 最新文章