深度学习-基于YOLOv5的NEU-DET钢材表面缺陷任务，加入CFPNet、动态卷积ODConv、多个检测头提升精度NEU-DET钢材表面缺陷 六类分别

为:‘crazing’,‘inclusion’，‘patches’,‘pitted_surface’‘rolled-in_scale’,'scratches

基于YOLOv5的NEU-DET钢材表面缺陷检测任务，并结合CFPNet、动态卷积ODConv和多个检测头来提升模型的精度。NEU-DET数据集包含六大类缺陷：‘crazing’, ‘inclusion’, ‘patches’, ‘pitted_surface’, ‘rolled-in_scale’, ‘scratches’。

1. 项目结构
   假设你的项目目录结构如下：

深色版本
neu_det_project/
├── dataset/
│ ├── images/
│ │ ├── train/
│ │ └── val/
│ ├── labels/
│ │ ├── train/
│ │ └── val/
├── models/
│ ├── yolov5s.pth
│ ├── yolov5m.pth
│ ├── yolov5l.pth
├── src/
│ ├── models/
│ │ ├── cfpnet.py
│ │ ├── odconv.py
│ │ ├── yolov5_custom.py
│ ├── utils/
│ │ ├── datasets.py
│ │ ├── general.py
│ │ ├── losses.py
│ │ ├── metrics.py
│ │ ├── plots.py
│ │ ├── torch_utils.py
│ ├── train.py
│ ├── predict.py
├── data.yaml
├── requirements.txt
2. 数据集配置文件 (data.yaml)
创建一个data.yaml文件，配置数据集的路径和类别信息：

yaml
深色版本
path: ./dataset # 数据集路径
train: images/train # 训练集图像路径
val: images/val # 验证集图像路径

nc: 6 # 类别数
names:

• crazing
• inclusion
• patches
• pitted_surface
• rolled-in_scale
• scratches

3. 安装依赖
   创建一个requirements.txt文件，列出所需的依赖项：

txt
深色版本
torch
torchvision
opencv-python
numpy
matplotlib
pyyaml
安装依赖项：

bash
深色版本
pip install -r requirements.txt
4. 模型定制
CFPNet
CFPNet（Contextual Feature Pyramid Network）是一种用于特征提取的网络，可以增强模型对不同尺度特征的捕捉能力。

python
深色版本

src/models/cfpnet.py

import torch
import torch.nn as nn

class CFPBlock(nn.Module):
def init(self, in_channels, out_channels):
super(CFPBlock, self).init()
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
self.relu = nn.ReLU(inplace=True)

def forward(self, x):
    x = self.conv1(x)
    x = self.relu(x)
    x = self.conv2(x)
    x = self.relu(x)
    return x

class CFPNet(nn.Module):
def init(self, num_classes):
super(CFPNet, self).init()
self.backbone = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2),
CFPBlock(64, 128),
nn.MaxPool2d(kernel_size=2, stride=2),
CFPBlock(128, 256),
nn.MaxPool2d(kernel_size=2, stride=2),
CFPBlock(256, 512),
nn.MaxPool2d(kernel_size=2, stride=2),
CFPBlock(512, 512)
)
self.classifier = nn.Sequential(
nn.Linear(512 * 7 * 7, 4096),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(4096, num_classes),
)

def forward(self, x):
    x = self.backbone(x)
    x = x.view(x.size(0), -1)
    x = self.classifier(x)
    return x

动态卷积ODConv
ODConv（Octave Convolution）是一种动态卷积方法，可以提高模型的效率和性能。

python
深色版本

src/models/odconv.py

import torch
import torch.nn as nn

class ODConv2d(nn.Module):
def init(self, in_channels, out_channels, kernel_size, stride=1, padding=0, groups=1, reduction=2):
super(ODConv2d, self).init()
self.in_channels = in_channels
self.out_channels = out_channels
self.kernel_size = kernel_size
self.stride = stride
self.padding = padding
self.groups = groups
self.reduction = reduction

    self.weight = nn.Parameter(torch.randn(out_channels, in_channels // groups, kernel_size, kernel_size))
    self.bias = nn.Parameter(torch.zeros(out_channels))

    self.fc = nn.Linear(in_channels // reduction, in_channels // groups)

def forward(self, x):
    b, c, h, w = x.size()
    weight = self.weight.unsqueeze(0).expand(b, -1, -1, -1, -1)
    weight = weight.contiguous().view(b * self.out_channels, c // self.groups, self.kernel_size, self.kernel_size)

    x = x.view(1, b * c, h, w)
    x = F.conv2d(x, weight, bias=self.bias.repeat(b), stride=self.stride, padding=self.padding, groups=b * self.groups)
    x = x.view(b, self.out_channels, x.size(2), x.size(3))

    return x

自定义YOLOv5模型
将CFPNet和ODConv集成到YOLOv5中，并添加多个检测头。

python
深色版本

src/models/yolov5_custom.py

from models.common import *
from models.experimental import *
from utils.general import make_divisible
from models.cfpnet import CFPNet
from models.odconv import ODConv2d

class YOLOv5Custom(nn.Module):
def init(self, cfg=‘yolov5s.yaml’, ch=3, nc=None, anchors=None):
super(YOLOv5Custom, self).init()
self.cfg = cfg
self.module_defs = parse_model_cfg(cfg)
self.hyperparams, self.module_list = create_modules(self.module_defs, ch, nc, anchors)
self.yolo_layers = get_yolo_layers(self)
self.cfpnet = CFPNet(nc)
self.odconv = ODConv2d(ch, ch, kernel_size=3, padding=1)

def forward(self, x):
    img_size = x.shape[-2:]
    yolo_out, out = [], []
    x = self.odconv(x)
    x = self.cfpnet(x)
    for i, module in enumerate(self.module_list):
        name = module.__class__.__name__
        if name in ['RouteLayer', 'ShortcutLayer']:
            x = module(x, out)
        elif name == 'YOLOLayer':
            yolo_out.append(module(x, img_size))
        else:
            x = module(x)
        out.append(x if self.routs[i] else [])
    return yolo_out if self.training else [torch.cat(x, 1) for x in zip(*yolo_out)]

5. 训练脚本 (train.py)
   python
   深色版本

src/train.py

import torch
import torch.optim as optim
import torch.nn.functional as F
from models.yolov5_custom import YOLOv5Custom
from utils.datasets import LoadImagesAndLabels
from utils.utils import non_max_suppression, scale_coords, xywh2xyxy, ap_per_class
from utils.losses import ComputeLoss

def train_model(data_yaml_path, model_config, epochs, batch_size, img_size, device):
# 加载数据集
dataset = LoadImagesAndLabels(data_yaml_path, img_size=img_size, batch_size=batch_size)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4)

# 加载模型
model = YOLOv5Custom(cfg=model_config, ch=3, nc=6).to(device)
model.train()

# 定义优化器和损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)
compute_loss = ComputeLoss(model)

for epoch in range(epochs):
    for i, (imgs, targets, _, _) in enumerate(dataloader):
        imgs = imgs.to(device)
        targets = targets.to(device)

        # 前向传播
        pred = model(imgs)

        # 计算损失
        loss, loss_items = compute_loss(pred, targets)

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if i % 10 == 0:
            print(f'Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(dataloader)}], Loss: {loss.item()}')

    # 保存模型
    torch.save(model.state_dict(), f'models/yolov5_custom_epoch_{epoch+1}.pth')

if name == “main”:
data_yaml_path = ‘data.yaml’
model_config = ‘models/yolov5s.yaml’
epochs = 100
batch_size = 16
img_size = 640
device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)

train_model(data_yaml_path, model_config, epochs, batch_size, img_size, device)

6. 预测脚本 (predict.py)
   python
   深色版本

src/predict.py

import cv2
import torch
from models.yolov5_custom import YOLOv5Custom
from utils.general import non_max_suppression, scale_coords, xywh2xyxy
from utils.datasets import letterbox

def predict_image(image_path, model_path, img_size=640, conf_thres=0.5, iou_thres=0.5):
# 加载模型
model = YOLOv5Custom(cfg=‘models/yolov5s.yaml’, ch=3, nc=6)
model.load_state_dict(torch.load(model_path))
model.eval()

# 读取图像
image = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_resized = letterbox(image_rgb, new_shape=img_size)[0]
image_resized = image_resized.transpose(2, 0, 1)
image_resized = torch.from_numpy(image_resized).float().div(255.0).unsqueeze(0)

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

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

# 处理预测结果
for det in pred:
    if det is not None and len(det):
        det[:, :4] = scale_coords(image_resized.shape[2:], det[:, :4], image.shape).round()
        for *xyxy, conf, cls in det:
            label = f'{["crazing", "inclusion", "patches", "pitted_surface", "rolled-in_scale", "scratches"][int(cls)]} {conf:.2f}'
            plot_one_box(xyxy, image, label=label, color=(0, 255, 0), line_thickness=2)

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

if name == “main”:
image_path = ‘path_to_your_image.jpg’
model_path = ‘models/yolov5_custom_epoch_100.pth’
predict_image(image_path, model_path)
7. 详细解释
数据集配置文件 (data.yaml)
path: 数据集的根目录路径。
train: 训练集图像的路径。
val: 验证集图像的路径。
nc: 类别数。
names: 类别名称列表。
训练脚本 (train.py)
导入依赖项：
import torch：导入PyTorch库。
import torch.optim as optim：导入优化器。
import torch.nn.functional as F：导入神经网络功能。
from models.yolov5_custom import YOLOv5Custom：导入自定义的YOLOv5模型。
from utils.datasets import LoadImagesAndLabels：导入数据加载器。
from utils.utils import non_max_suppression, scale_coords, xywh2xyxy, ap_per_class：导入辅助函数。
from utils.losses import ComputeLoss：导入损失函数。
定义训练函数：
train_model：加载数据集和模型，设置优化器和损失函数，进行训练，并保存模型。
主函数：
设置数据集路径、模型配置、训练参数等。
调用train_model函数进行训练。
预测脚本 (predict.py)
导入依赖项：
import cv2：导入OpenCV库。
import torch：导入PyTorch库。
from models.yolov5_custom import YOLOv5Custom：导入自定义的YOLOv5模型。
from utils.general import non_max_suppression, scale_coords, xywh2xyxy：导入辅助函数。
from utils.datasets import letterbox：导入图像预处理函数。
定义预测函数：
predict_image：加载模型，读取图像，进行预测，处理预测结果，并显示带有标注的图像。
主函数：
设置图像路径和模型路径。
调用predict_image函数进行预测。
8. 注意事项
数据集路径：确保数据集路径正确，特别是data.yaml文件中的路径。
模型配置：确保模型配置文件路径正确。
图像大小：img_size可以根据实际需求调整，通常使用640或1280。
设备：确保设备（CPU或GPU）可用。
9. 数据增强
为了增加数据集的多样性，可以使用数据增强技术。YOLOv5在训练过程中默认支持多种数据增强方法，如随机裁剪、翻转、颜色抖动等。如果需要自定义数据增强，可以参考YOLOv5的文档进行配置。

下面是训练数据集时的关键代码，包括数据加载、模型初始化、训练循环和时间记录。这些代码片段将帮助你完成基于YOLOv5的NEU-DET钢材表面缺陷检测任务的训练。

1. 数据加载

首先，我们需要加载数据集。这里使用了LoadImagesAndLabels类来加载图像和标签。

# src/utils/datasets.py
import os
import random
import cv2
import numpy as np
from torch.utils.data import Dataset

class LoadImagesAndLabels(Dataset):
    def __init__(self, data_yaml_path, img_size=640, batch_size=16, augment=False):
        with open(data_yaml_path, 'r') as f:
            data = yaml.safe_load(f)
        
        self.img_size = img_size
        self.batch_size = batch_size
        self.augment = augment
        self.img_files = []
        self.label_files = []

        for d in ['train', 'val']:
            path = os.path.join(data['path'], data[d])
            for file in os.listdir(path):
                if file.endswith('.jpg'):
                    self.img_files.append(os.path.join(path, file))
                    self.label_files.append(os.path.join(path.replace('images', 'labels'), file.replace('.jpg', '.txt')))

    def __len__(self):
        return len(self.img_files)

    def __getitem__(self, index):
        img_path = self.img_files[index]
        label_path = self.label_files[index]

        # Load image
        img = cv2.imread(img_path)
        img = cv2.resize(img, (self.img_size, self.img_size))
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
        img = np.ascontiguousarray(img, dtype=np.float32) / 255.0

        # Load labels
        if os.path.exists(label_path):
            with open(label_path, 'r') as f:
                labels = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32)
        else:
            labels = np.zeros((0, 5), dtype=np.float32)

        return torch.from_numpy(img), torch.from_numpy(labels), img_path, (self.img_size, self.img_size)

2. 模型初始化

接下来，我们需要初始化模型。这里使用了自定义的YOLOv5Custom模型，集成了CFPNet和ODConv。

# src/models/yolov5_custom.py
import torch
import torch.nn as nn
from models.common import *
from models.experimental import *
from utils.general import make_divisible
from models.cfpnet import CFPNet
from models.odconv import ODConv2d

class YOLOv5Custom(nn.Module):
    def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None):
        super(YOLOv5Custom, self).__init__()
        self.cfg = cfg
        self.module_defs = parse_model_cfg(cfg)
        self.hyperparams, self.module_list = create_modules(self.module_defs, ch, nc, anchors)
        self.yolo_layers = get_yolo_layers(self)
        self.cfpnet = CFPNet(nc)
        self.odconv = ODConv2d(ch, ch, kernel_size=3, padding=1)

    def forward(self, x):
        img_size = x.shape[-2:]
        yolo_out, out = [], []
        x = self.odconv(x)
        x = self.cfpnet(x)
        for i, module in enumerate(self.module_list):
            name = module.__class__.__name__
            if name in ['RouteLayer', 'ShortcutLayer']:
                x = module(x, out)
            elif name == 'YOLOLayer':
                yolo_out.append(module(x, img_size))
            else:
                x = module(x)
            out.append(x if self.routs[i] else [])
        return yolo_out if self.training else [torch.cat(x, 1) for x in zip(*yolo_out)]

3. 训练循环

在训练循环中，我们记录每个epoch的训练时间和总训练时间。

# src/train.py
import torch
import torch.optim as optim
import torch.nn.functional as F
from models.yolov5_custom import YOLOv5Custom
from utils.datasets import LoadImagesAndLabels
from utils.utils import non_max_suppression, scale_coords, xywh2xyxy, ap_per_class
from utils.losses import ComputeLoss
import time
import yaml

def train_model(data_yaml_path, model_config, epochs, batch_size, img_size, device):
    # 加载数据集
    dataset = LoadImagesAndLabels(data_yaml_path, img_size=img_size, batch_size=batch_size)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4)

    # 加载模型
    model = YOLOv5Custom(cfg=model_config, ch=3, nc=6).to(device)
    model.train()

    # 定义优化器和损失函数
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    compute_loss = ComputeLoss(model)

    # 记录训练时间
    start_time = time.time()

    for epoch in range(epochs):
        epoch_start_time = time.time()
        for i, (imgs, targets, _, _) in enumerate(dataloader):
            imgs = imgs.to(device)
            targets = targets.to(device)

            # 前向传播
            pred = model(imgs)

            # 计算损失
            loss, loss_items = compute_loss(pred, targets)

            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if i % 10 == 0:
                print(f'Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(dataloader)}], Loss: {loss.item()}')

        epoch_end_time = time.time()
        epoch_duration = epoch_end_time - epoch_start_time
        print(f'Epoch [{epoch+1}/{epochs}] completed in {epoch_duration:.2f} seconds')

        # 保存模型
        torch.save(model.state_dict(), f'models/yolov5_custom_epoch_{epoch+1}.pth')

    end_time = time.time()
    total_duration = end_time - start_time
    print(f'Total training time: {total_duration:.2f} seconds')

if __name__ == "__main__":
    data_yaml_path = 'data.yaml'
    model_config = 'models/yolov5s.yaml'
    epochs = 100
    batch_size = 16
    img_size = 640
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    train_model(data_yaml_path, model_config, epochs, batch_size, img_size, device)

4. 详细解释

数据加载

• LoadImagesAndLabels类从data.yaml文件中读取数据集路径，并加载图像和标签。
• 图像被缩放到指定的大小，并转换为RGB格式和归一化。
• 标签文件被读取并转换为NumPy数组。

模型初始化

• YOLOv5Custom类继承自nn.Module，并集成了CFPNet和ODConv。
• 模型的前向传播方法中，先使用ODConv处理输入图像，再通过CFPNet提取特征，最后通过YOLOv5的模块列表进行前向传播。

训练循环

• 训练循环中，记录每个epoch的开始和结束时间，计算每个epoch的训练时间。
• 在每个step中，进行前向传播、计算损失、反向传播和优化器更新。
• 每10个step打印一次当前的损失值。
• 每个epoch结束后，保存当前模型的权重。
• 记录整个训练过程的总时间，并在训练结束后打印。

5. 运行脚本

1. 训练模型：

   python src/train.py
   

6. 注意事项

1. 数据集路径：确保数据集路径正确，特别是data.yaml文件中的路径。
2. 模型配置：确保模型配置文件路径正确。
3. 图像大小：img_size可以根据实际需求调整，通常使用640或1280。
4. 设备：确保设备（CPU或GPU）可用。

7. 总结

通过以上步骤，在训练脚本中记录每个epoch的训练时间和整个训练过程的总时间。这将帮助你更好地了解训练过程的性能和效率