PCB缺陷检测 - YOLOv8实现

基于YOLOv8的PCB（印刷电路板）缺陷检测系统，支持模型训练、微调和验证。

项目概述

本项目使用YOLOv8目标检测模型来识别和定位PCB板上的缺陷。项目包含完整的训练、微调和验证流程，适用于PCB制造质量控制场景。

环境要求

• Python 3.8+
• PyTorch 2.0+（支持CUDA）
• CUDA Toolkit 11.8+（用于GPU加速）
• Ultralytics 8.0+

安装指南

1. 克隆仓库：

   git clone <repository-url>
   cd YOLO

2. 安装依赖：

   pip install -r requirements.txt

3. 验证GPU可用性：

   python test_gpu.py

项目结构

.
├── main_pcb.py              # 主训练脚本
├── main_pcb_strengthen.py   # 模型微调脚本
├── main_pcb_validate.py     # 模型验证脚本
├── test_gpu.py              # GPU测试工具
├── requirements.txt         # 项目依赖
├── yolov8n.pt               # YOLOv8预训练模型
├── try.v1i.yolov8/          # 数据集目录
│   └── data.yaml            # 数据集配置文件
├── runs/                    # 训练结果输出目录
└── myself_vali/             # 自定义验证数据

使用指南

1. 模型训练

运行以下命令开始训练PCB缺陷检测模型：

python main_pcb.py

训练参数可在脚本中调整，包括：

• 训练轮数（epochs）
• 批处理大小（batch size）
• 图像尺寸（image size）
• 输出目录名称

2. 模型微调

对已训练的模型进行微调，提高性能：

python main_pcb_strengthen.py --base_model runs/detect/pcb_demo/weights/best.pt

3. 模型验证

评估模型性能：

python main_pcb_validate.py --model_path runs/detect/pcb_demo_strengthen/weights/best.pt

数据集准备

项目使用YOLO格式的数据集，需要包含以下结构：

try.v1i.yolov8/
├── data.yaml    # 数据集配置
├── train/       # 训练图像和标签
│   ├── images/
│   └── labels/
└── val/         # 验证图像和标签
    ├── images/
    └── labels/

data.yaml 文件应包含类别信息和路径配置。

模型输出

训练和验证结果保存在 runs/ 目录下：

• runs/detect/ - 训练结果
• runs/val/ - 验证结果

每次运行会创建一个新的子目录，包含：

• 训练好的模型权重（weights/best.pt）
• 训练过程指标图表
• 验证结果和性能指标

性能指标

模型评估使用以下指标：

• mAP50-95：在不同IoU阈值下的平均精度
• mAP50：在IoU=0.5时的平均精度
• 精确率（Precision）和召回率（Recall）

注意事项

• 确保CUDA环境正确配置以启用GPU加速
• 对于大型数据集，可能需要调整batch size以适应GPU内存
• 微调时使用较小的学习率（默认0.0001）

许可证

[指定许可证类型]

联系方式

[您的联系信息]