如何训练使用DeepPCB数据集（1500张）六种缺陷类型：并实现可视化及评估 YOLOv5算法实现对PCB板上的缺陷进行检测

DeepPCB数据集（1500张）
可用于yolov5，v8等模型训练

提供YOLOv5算法实现对PCB板上的缺陷进行检测识别代码。

标注了六种缺陷类型：open、short、mousebite、spur、pin-hole、spur。

YOLOv5 PCB缺陷检测

import os
import glob
import yaml
from pathlib import Path
import torch
from IPython.display import Image, clear_output
from ultralytics.yolo.v5 import train

# 设置随机种子以保证可重复性
torch.manual_seed(42)

# 定义数据集路径
dataset_dir = 'path/to/dataset'

# 创建YOLOv5的数据集配置文件
data_config = {
    'train': os.path.join(dataset_dir, 'train/images'),
    'val': os.path.join(dataset_dir, 'val/images'),
    'test': os.path.join(dataset_dir, 'test/images'),
    'nc': 5,  # 类别数量
    'names': ['open', 'short', 'mousebite', 'spur', 'pin-hole']  # 类别名称
}

with open(os.path.join(dataset_dir, 'data.yaml'), 'w') as f:
    yaml.dump(data_config, f)

# 打印数据配置信息
print("Data Configuration:")
print(data_config)

# 训练模型
results = train.run(
    data=os.path.join(dataset_dir, 'data.yaml'),
    imgsz=640,
    batch=16,
    epochs=100,
    weights='yolov5s.pt',
    name='pcb_defect_detection',
    project='runs/train'
)

# 评估模型
# YOLOv5的评估通常在训练过程中自动完成，但可以手动运行
# results = val.run(weights='runs/train/pcb_defect_detection/weights/best.pt', data=data_config)

# 可视化预测结果
source_image = '../path/to/dataset/val/sample.jpg'  # 替换为你要测试的图片路径
results = detect.run(
    source=source_image,
    weights='runs/train/pcb_defect_detection/weights/best.pt',
    conf_thres=0.25,
    iou_thres=0.45,
    save=True,
    save_txt=True
)

# 显示预测结果
Image(filename='runs/detect/exp/sample.jpg')



假设
你有一个DeepPCB数据集，包含1500张图像，并且标注了六种缺陷类型：`open`、`short`、`mousebite`、`spur`、`pin-hole` 和 `spur`（注意：这里有两个 `spur` 类别，可能是重复的，假设只保留一个）。数据集可以直接用于YOLOv5和YOLOv8等模型训练。

### 项目介绍

#### 数据准备
- **数据集**: 包含1500张图像。
- **类别**:
  - `open`: 开路
  - `short`: 短路
  - `mousebite`: 鼠咬痕
  - `spur`: 刺点
  - `pin-hole`: 孔洞
- **标注格式**: YOLO TXT格式。
- **数据分布**:
  - 假设数据集已经按照一定的比例划分好了训练集、验证集和测试集。

#### 模型选择
- **YOLOv5**: 使用YOLOv5进行目标检测。YOLOv5是YOLO系列的一个流行版本，具有较高的性能和良好的精度。

#### 功能
- **数据加载**: 自动从指定目录加载图像和标注文件。
- **模型训练**: 使用YOLOv5进行训练。
- **模型评估**: 在验证集上评估模型性能。
- **结果保存**: 保存训练日志和最佳模型权重。
- **可视化预测结果**: 可视化预测结果以进行验证。

### 代码实现

首先，确保你已经安装了YOLOv5库和其他必要的依赖项。你可以通过以下命令安装YOLOv5：

```bash
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

接下来，我们编写代码来组织数据集并训练YOLOv5模型。假设数据集已经是标准的YOLO格式，并且按照给定的比例划分好了训练集、验证集和测试集。

如何使用这些代码

1. 准备数据：

   • 确保你的数据集格式正确，包含训练集、验证集和测试集文件夹，以及对应的标注文件。
   • 示例数据结构如下：

     path/to/dataset/
     ├── train/
     │   ├── images/
     │   │   ├── image1.jpg
     │   │   ├── image2.jpg
     │   │   └── ...
     │   ├── labels/
     │   │   ├── image1.txt
     │   │   ├── image2.txt
     │   │   └── ...
     ├── val/
     │   ├── images/
     │   │   ├── image1.jpg
     │   │   ├── image2.jpg
     │   │   └── ...
     │   ├── labels/
     │   │   ├── image1.txt
     │   │   ├── image2.txt
     │   │   └── ...
     ├── test/
     │   ├── images/
     │   │   ├── image1.jpg
     │   │   ├── image2.jpg
     │   │   └── ...
     │   ├── labels/
     │   │   ├── image1.txt
     │   │   ├── image2.txt
     │   │   └── ...
     ├── data.yaml
     

2. 替换数据路径：

   • 在代码中，将 'path/to/dataset' 替换为你的数据集路径。

   dataset_dir = 'your_dataset_directory'
   

3. 运行代码：

   • 将上述代码复制到你的Python脚本中，并运行该脚本。
   • 确保你已经安装了所需的库：

     git clone https://github.com/ultralytics/yolov5
     cd yolov5
     pip install -r requirements.txt
     

示例：使用自定义数据集

假设你有一个新的数据集 my_deep_pcb_dataset，其内容如下：

my_deep_pcb_dataset/
├── train/
│   ├── images/
│   │   ├── image1.jpg
│   │   ├── image2.jpg
│   │   └── ...
│   ├── labels/
│   │   ├── image1.txt
│   │   ├── image2.txt
│   │   └── ...
├── val/
│   ├── images/
│   │   ├── image1.jpg
│   │   ├── image2.jpg
│   │   └── ...
│   ├── labels/
│   │   ├── image1.txt
│   │   ├── image2.txt
│   │   └── ...
├── test/
│   ├── images/
│   │   ├── image1.jpg
│   │   ├── image2.jpg
│   │   └── ...
│   ├── labels/
│   │   ├── image1.txt
│   │   ├── image2.txt
│   │   └── ...
├── data.yaml

你可以按照以下步骤进行替换：

1. 修改数据路径：

   dataset_dir = 'my_deep_pcb_dataset'
   

2. 运行完整的代码：

   • 将所有代码整合到一个Python脚本中，并运行该脚本。

注释说明

代码中包含了详细的注释，帮助你理解每个部分的功能。以下是关键部分的注释：

• 数据准备：

  • data_config: 定义训练集、验证集和测试集的路径，以及类别信息。

• 模型训练：

  • train.run: 使用YOLOv5进行训练。

• 模型评估：

  • YOLOv5的评估通常在训练过程中自动完成，但可以手动运行 val.run 进行评估。

• 可视化预测结果：

  • detect.run: 进行推理并显示预测结果。

结果

运行代码后，你将得到以下结果：

1. 控制台输出：

   • 训练过程中每个epoch的日志信息。
   • 验证集上的评价指标（如mAP）。

2. 文件输出：

   • runs/train/pcb_defect_detection/weights/best.pt: 最佳模型权重。
   • runs/val/exp/results.txt: 验证结果。

3. 图像输出：

   • runs/detect/exp/sample.jpg: 带有预测边界的图像。

希望这些详细的信息和代码能够帮助你顺利实施和优化你的项目

运行步骤总结

1. 克隆YOLOv5仓库：

   git clone https://github.com/ultralytics/yolov5
   cd yolov5
   

2. 安装依赖项：

   pip install -r requirements.txt
   

3. 准备数据集：

   • 确保数据集路径正确，并且包含训练集、验证集和测试集文件夹，以及对应的标注文件。

4. 运行训练脚本：

   python your_script_name.py
   

5. 评估模型：

   • 评估结果会在训练结束后自动输出在控制台和文件中。

6. 可视化预测结果：

   • 测试图像的结果会保存在 runs/detect/exp/ 文件夹中，可以直接查看带有预测边界的图像。

希望这些详细的指导和代码示例能帮助你成功实现和优化你的PCB缺陷检测项目。