基于改进YOLOv8算法的钢材瑕疵辅助检测系统
传统钢材缺陷检测主要依赖于人工视觉检查或使用简单的机器视觉系统,但这些方法存在效率低下、准确度不高、劳动强度大以及易受人为因素影响等问题。因此,我独立开发了一种基于改进YOLOv8算法的钢材缺陷辅助检测系统。
方法
项目提出了一个改进版的YOLOv8算法。首先,提出了一种残差注意力机制——ResCBAM,它结合了残差连接,并将其添加到YOLOv8模型的头部以提升模型提取特征的能力同时防止过拟合。GhostConv被用来替换骨干网络中的普通卷积,这使得模型在保持性能的同时更加轻量化。
同时,引入WIOU损失函数替代原有的CIOU损失函数,避免了小目标对损失函数产生过大影响的问题,缓解了数据集不平衡问题。实验结果显示,相较于原始YOLOv8模型,在公共NEU-DET数据集上提出的改进YOLOv8模型的精确率(Precision)、F1得分(F1score)和mAP50分别提高了5.6%、3.1%和1.5%,mAP50达到了79.2%,显著提升了检测的准确性和效率,对于保证产品质量和降低检测成本具有重要意义。该模型还部署在WEB端,用户可以通过访问网页上传图片或拍照进行检测,并可以保存结果。
架构细节
ResCBAM模块
模型 | 测试尺寸 | 参数量 | FLOPs | mAP50val | mAP50-95val |
---|---|---|---|---|---|
YOLOv5s | 640 | 7.04M | 15.9G | 70.87% | 35.02% |
YOLOv5m | 640 | 20.89M | 183.5G | 71.74% | 36.67% |
YOLOv5l | 640 | 46.17M | 183.5G | 72.23% | 36.74% |
YOLOv7 | 640 | 37.22M | 196.2G | 71.92% | 37.2% |
YOLOv8 | 640 | 3.01M | 8.1G | 77.7% | 46.2% |
我们的模型 | 640 | 4.05M | 10.2G | 79.2% | 47% |
通过消融研究进一步验证了各个组件的有效性:
模型 | 测试尺寸 | 参数量 | FLOPs | mAP50val |
---|---|---|---|---|
YOLOv8 | 640 | 3.01M | 8.1G | 77.7% |
YOLOv8+GhostConv | 640 | 2.8M | 7.8G | 78.7% |
YOLOv8+ResCBAM | 640 | 4.24M | 10.5G | 78.2% |
YOLOv8+ResCBAM+GhostConv(我们的模型) | 640 | 4.05M | 10.2G | 79.2% |
此外,还包括PR曲线、混淆矩阵等模型指标展示,以全面评估模型性能。
要求
- 系统:Linux (Ubuntu)/Windows
- Python版本:>= 3.8
- Pytorch版本:>= 1.13.1
- NVIDIA GPU + CUDA CuDNN
数据集
可以从以下链接下载NEU-DET数据集:。此数据集已按80-10-10%的比例划分为训练集、验证集和测试集。
快速开始
- 安装:
pip install -r requirements.txt
-
下载并解压数据集,将相关文件夹移动到项目目录下。
-
训练与验证命令示例:
python train.py --model .yolov8_ResBlock_CBAM_GhostConv.yaml --data_dir NEU-DET/data.ymal
本项目的开发不仅推动了钢材缺陷检测技术的发展,也为其他工业领域的质量控制提供了新的思路和方法。