如何使用YOLOv8训练芯片缺陷数据集包括数据准备、模型训练、评估和结果可视化等多个步骤 558张，标签为yolo格式 4类缺陷类型（自定义）：划痕、污渍、封装破裂、引脚-CSDN博客

在这里插入图片描述
芯片缺陷数据集，可用于深度学习模型训练
为本人制作实物并拍摄，未经过数据增强
共558张，标签为yolo格式
缺陷类型（自定义）：划痕、污渍、封装破裂、引脚
用于yolov8模型训练

558张芯片缺陷图像的数据集，标签格式为YOLO格式。将从数据准备、模型训练、评估和可视化等方面详细介绍整个流程。使用YOLOv8进行训练。

1. 环境准备

首先，确保你已经安装了必要的库和工具。你可以使用以下命令安装所需的库：

pip install ultralytics
pip install torch torchvision
pip install opencv-python
pip install pandas
pip install matplotlib

2. 数据准备

假设你的数据集目录结构如下：

dataset/
├── images/
│   ├── train/
│   └── val/
└── labels/
    ├── train/
    └── val/

其中，images文件夹包含训练和验证的图像，labels文件夹包含相应的YOLO格式标签文件。

3. 随机划分训练集与验证集

我们将编写一个脚本来随机划分训练集和验证集，并将相应的图像和标签文件移动到相应的文件夹中。

import os
import random
import shutil

def split_dataset(image_dir, label_dir, train_ratio=0.8):
    images = os.listdir(image_dir)
    random.shuffle(images)

    train_size = int(len(images) * train_ratio)
    train_images = images[:train_size]
    val_images = images[train_size:]

    train_image_dir = os.path.join(image_dir, 'train')
    val_image_dir = os.path.join(image_dir, 'val')
    train_label_dir = os.path.join(label_dir, 'train')
    val_label_dir = os.path.join(label_dir, 'val')

    os.makedirs(train_image_dir, exist_ok=True)
    os.makedirs(val_image_dir, exist_ok=True)
    os.makedirs(train_label_dir, exist_ok=True)
    os.makedirs(val_label_dir, exist_ok=True)

    for image in train_images:
        shutil.move(os.path.join(image_dir, image), os.path.join(train_image_dir, image))
        label = os.path.splitext(image)[0] + '.txt'
        shutil.move(os.path.join(label_dir, label), os.path.join(train_label_dir, label))

    for image in val_images:
        shutil.move(os.path.join(image_dir, image), os.path.join(val_image_dir, image))
        label = os.path.splitext(image)[0] + '.txt'
        shutil.move(os.path.join(label_dir, label), os.path.join(val_label_dir, label))

# 数据集路径
image_dir = 'dataset/images'
label_dir = 'dataset/labels'

# 调用函数
split_dataset(image_dir, label_dir)

4. 创建数据配置文件

YOLOv8需要一个数据配置文件来指定数据集的路径和其他相关信息。创建一个名为chip_defects.yaml的文件，内容如下：

# Chip Defects Dataset Configuration

# Path to the dataset directory
path: ./dataset

# Training and validation image directories
train: images/train
val: images/val

# Number of classes
nc: 4

# Class names
names:
  0: 划痕
  1: 污渍
  2: 封装破裂
  3: 引脚

5. 训练模型

使用YOLOv8进行训练非常简单。你可以使用以下命令来启动训练：

yolo train data=chip_defects.yaml model=yolov8n.pt epochs=100 imgsz=640

解释：

data=chip_defects.yaml: 指定数据配置文件。
model=yolov8n.pt: 使用预训练的YOLOv8小模型（yolov8n）。你可以选择其他大小的模型，如yolov8s、yolov8m、yolov8l或yolov8x。
epochs=100: 训练的轮数。
imgsz=640: 图像的尺寸。

6. 评估模型

训练完成后，你可以使用以下命令来评估模型在验证集上的性能：

yolo val data=chip_defects.yaml model=runs/detect/train/weights/best.pt imgsz=640

解释：

data=chip_defects.yaml: 指定数据配置文件。
model=runs/detect/train/weights/best.pt: 指定训练过程中保存的最佳模型权重文件。
imgsz=640: 图像的尺寸。

7. 可视化预测结果

你可以使用以下Python代码来可视化模型的预测结果：

import cv2
import torch
from ultralytics import YOLO

# 加载模型
model = YOLO('runs/detect/train/weights/best.pt')

# 读取图像
image_path = 'dataset/images/val/your_image.jpg'
image = cv2.imread(image_path)

# 进行预测
results = model(image)

# 可视化预测结果
for result in results:
    boxes = result.boxes.xyxy
    confidences = result.boxes.conf
    class_ids = result.boxes.cls

    for box, conf, class_id in zip(boxes, confidences, class_ids):
        x1, y1, x2, y2 = map(int, box)
        label = model.names[int(class_id)]
        confidence = float(conf)

        # 绘制边界框
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        text = f'{label}: {confidence:.2f}'
        cv2.putText(image, text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

# 显示图像
cv2.imshow('Prediction', image)
cv2.waitKey(0)
cv2.destroyAllWindows()