基于YOLOv8的金属品质检测系统开发与实践

1. 项目概述

金属品质检测是工业生产中至关重要的环节，直接影响产品的安全性和可靠性。传统的人工检测方法效率低下且容易出错，而基于深度学习的自动化检测系统正在改变这一现状。本项目基于YOLOv8构建了一套完整的金属品质检测系统，包含数据集标注、模型训练、性能优化和Web前端展示的全流程解决方案。

1.1 核心功能特点

• 高效检测架构：采用改进版YOLOv8模型，在保持实时性的同时提升检测精度
• 完整数据支持：提供3714张已标注的金属表面图像数据集（Good/Bad两类）
• 端到端解决方案：从数据准备到模型部署的一站式实现
• 可视化界面：基于Streamlit开发的Web展示系统
• 工业级优化：针对实际生产环境进行多维度性能调优

提示：系统默认使用GPU加速，但也可在纯CPU环境下运行（速度会有所下降）

2. 系统架构设计

2.1 技术栈组成

2.2 模型改进方案

原始YOLOv8模型在金属检测场景下进行了以下关键改进：

1. 注意力机制增强：

   • 在Backbone末端添加CBAM注意力模块
   • 使用BiFPN替代原FPN结构

   class CBAM(nn.Module): def __init__(self, channels, reduction=16): super().__init__() self.channel_attention = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(channels, channels//reduction, 1), nn.ReLU(), nn.Conv2d(channels//reduction, channels, 1), nn.Sigmoid() ) self.spatial_attention = nn.Sequential( nn.Conv2d(2, 1, 7, padding=3), nn.Sigmoid() )

2. 损失函数优化：

   • 使用Focal Loss替代原分类损失
   • 引入SIoU作为回归损失

   def siou_loss(pred, target): # SIoU计算实现 ch = torch.pow(pred[:,:2] - target[:,:2], 2).sum(dim=1) # ...完整计算过程... return 1 - siou

3. 数据增强策略：

   • 针对金属反光特性的Mosaic增强
   • 随机HSV调整（降低饱和度变化幅度）

3. 数据集构建与处理

3.1 数据集规格

• 总样本量：3714张工业现场采集图像
• 类别分布：
  • Good-Metal：2486张（合格金属）
  • Bad-Metal：1228张（缺陷金属）
• 缺陷类型覆盖：
  • 表面划痕（35%）
  • 锈蚀斑点（28%）
  • 结构变形（22%）
  • 其他缺陷（15%）

3.2 标注规范示例

数据集采用YOLO格式标注，每个图像对应一个.txt标注文件：

# 示例标注内容 0 0.543 0.712 0.125 0.156 # 类别 x_center y_center width height 1 0.231 0.456 0.078 0.112

关键处理脚本：

def convert_annotation(image_path): """将原始标注转换为YOLO格式""" img = cv2.imread(image_path) h, w = img.shape[:2] with open(annotation_path) as f: # 解析原始标注并转换坐标 x_min, y_min, x_max, y_max = parse_raw_annotation(f.read()) x_center = (x_min + x_max) / 2 / w y_center = (y_min + y_max) / 2 / h width = (x_max - x_min) / w height = (y_max - y_min) / h return f"{class_id} {x_center} {y_center} {width} {height}"

3.3 数据增强策略

针对金属检测的特殊性，采用以下增强组合：

train_transform = A.Compose([ A.RandomBrightnessContrast(p=0.5), A.GaussNoise(var_limit=(10, 50), p=0.3), A.Rotate(limit=15, p=0.5), A.HorizontalFlip(p=0.5), A.Cutout(max_h_size=20, max_w_size=20, p=0.3), ], bbox_params=A.BboxParams(format='yolo'))

4. 模型训练与优化

4.1 训练配置

关键训练参数设置：

# data/metal.yaml train: ../datasets/Metal-Detection/train/images val: ../datasets/Metal-Detection/valid/images nc: 2 # 类别数 names: ['Good-Metal', 'Bad-Metal']

启动训练命令：

python train.py --img 640 --batch 16 --epochs 100 --data data/metal.yaml \ --weights yolov8s.pt --device 0 --name metal_detection

4.2 关键训练技巧

1. 学习率调度：

   • 采用余弦退火策略
   • 初始lr=0.01，最终lr=0.0001

   lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - lrf) + lrf

2. 早停策略：

   • 监控验证集mAP@0.5
   • patience=20（20个epoch无提升则停止）

3. 混合精度训练：

   scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): pred = model(imgs) loss = compute_loss(pred, targets) scaler.scale(loss).backward() scaler.step(optimizer)

4.3 性能指标

在测试集上的表现：

5. 系统部署与使用

5.1 Web界面开发

基于Streamlit的检测界面核心代码：

def main(): st.title("金属品质检测系统") uploaded_file = st.file_uploader("上传金属图像", type=['jpg','png']) if uploaded_file: img = Image.open(uploaded_file) st.image(img, caption='上传图像', use_column_width=True) if st.button('开始检测'): with st.spinner('检测中...'): results = detect(img) plot = plot_results(results) st.image(plot, caption='检测结果', use_column_width=True)

5.2 生产环境部署

Docker部署方案：

FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime WORKDIR /app COPY requirements.txt . RUN pip install -r requirements.txt COPY . . EXPOSE 8501 CMD ["streamlit", "run", "web.py", "--server.port=8501"]

启动命令：

docker build -t metal-detection . docker run -p 8501:8501 --gpus all metal-detection

5.3 性能优化技巧

1. TensorRT加速：

   model = YOLO('best.pt') model.export(format='engine', device=0)

2. 多线程处理：

   from concurrent.futures import ThreadPoolExecutor def batch_detect(images): with ThreadPoolExecutor(max_workers=4) as executor: results = list(executor.map(detect, images)) return results

3. 缓存机制：

   @st.cache_resource def load_model(): return YOLO('best.pt')

6. 常见问题解决

6.1 训练问题排查

6.2 部署问题处理

1. CUDA内存不足：

   export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

2. Streamlit响应慢：

   • 启用页面缓存
   • 减少实时检测的图像分辨率

3. 模型加载失败：

   try: model = torch.load('best.pt', map_location='cpu') except RuntimeError: model = YOLO('best.pt') # 使用ultralytics原生加载

在实际部署中，我们发现金属反光问题会显著影响检测效果。通过添加以下预处理步骤可提升稳定性：

def reduce_glare(image): lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB) l, a, b = cv2.split(lab) clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8)) limg = cv2.merge([clahe.apply(l), a, b]) return cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)