基于PyTorch和YOLOv5的完整代码训练使用——衣物皮革类数据集/织物布匹缺陷检测数据集 5921张(6类）深度学习目标检测模型训练

好的，织物疵点检测是一个重要的应用领域，可以显著提高纺织品的质量控制和生产效率。
以下是一个基于PyTorch和YOLOv5的完整代码示例，涵盖了数据准备、模型训练、验证和评估等关键步骤。我们将使用YOLOv5作为目标检测模型，因为它在实时检测任务中表现出色。在这里插入图片描述

1. 环境准备

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

pip install torch torchvision opencv-python pandas
pip install -U albumentations
pip install -U pyyaml
pip install -U yolov5

在这里插入图片描述

2. 数据准备

假设你的数据集目录结构如下：
在这里插入图片描述

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

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

3. 创建数据配置文件

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

# Fabric 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: 6

# Class names
names:
  0: 断经
  1: 断纬
  2: 缩纬
  3: 缩经
  4: 擦伤
  5: 污渍

4. 训练模型

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

yolo train data=fabric_defects.yaml model=yolov5s.pt epochs=100 imgsz=640

解释：

data=fabric_defects.yaml: 指定数据配置文件。
model=yolov5s.pt: 使用预训练的YOLOv5小模型（yolov5s）。你可以选择其他大小的模型，如yolov5m、yolov5l或yolov5x。
epochs=100: 训练的轮数。
imgsz=640: 图像的尺寸。

5. 评估模型

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

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

解释：

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

6. 可视化预测结果

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

import cv2
import torch
from yolov5.models.experimental import attempt_load
from yolov5.utils.general import non_max_suppression, scale_coords
from yolov5.utils.plots import Annotator

# 加载模型
model = attempt_load('runs/detect/train/weights/best.pt', map_location='cuda' if torch.cuda.is_available() else 'cpu')
model.eval()

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

# 图像预处理
img = cv2.resize(image, (640, 640))
img = img.transpose(2, 0, 1)  # HWC to CHW
img = img.astype(np.float32) / 255.0
img = np.expand_dims(img, axis=0)
img = torch.from_numpy(img).to('cuda' if torch.cuda.is_available() else 'cpu')

# 进行预测
with torch.no_grad():
    pred = model(img)[0]

# 非极大值抑制
pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)

# 可视化预测结果
annotator = Annotator(image.copy(), line_width=2, example='断经 断纬 缩纬 缩经 擦伤 污渍')
for det in pred:
    if det is not None and len(det):
        det[:, :4] = scale_coords(img.shape[2:], det[:, :4], image.shape).round()
        for *xyxy, conf, cls in reversed(det):
            c = int(cls)
            label = f'{model.names[c]} {conf:.2f}'
            annotator.box_label(xyxy, label, color=(0, 255, 0))

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