Hello，大家好，我是augustqi。

今天给大家分享一个图像分类项目：基于yolov5的钢板表面缺陷分类(附代码和数据集)，在之前的文章中，我们使用yolov7算法对钢板表面缺陷进行检测，属于目标检测范畴，鉴于yolov5-6.2和yolov5-7.0版本中增加了图像分类功能，我们尝试使用yolov5做一个图像分类任务，使用的是钢板表面缺陷数据集。

钢板表面缺陷检测文章回顾：
【表面缺陷检测】基于yolov7的钢板表面缺陷检测(Ubuntu系统)
【表面缺陷检测】基于yolov7的钢板表面缺陷检测(附代码)

本项目也是从零开始制作自己的数据集，学会如何使用yolov5训练图像分类模型的保姆级教程，多的不说，少的不唠，下面我们一起来学习一下吧。

以下内容，完全是我根据参考资料和个人理解撰写出来的，不存在滥用原创的问题。

1. 引言

从事深度学习行业的都知道，yolov5是一个很优秀的目标检测框架，代码已开源至github，截止到2023年1月14日，获得了34.5k Star，12.5k Fork，妥妥的人类高质量项目。项目也在持续更新，并不断完善和增加了新的功能，例如在yolov5-6.2和yolov5-7.0版本中增加了图像分类功能，yolov5-7.0版本中增加了图像分割功能。

资料显示，官方使用4张A100显卡，在ImageNet数据集上训练了90个epoch，得到了YOLOv5-cls分类模型，同时训练了ResNet和EfficientNet模型进行比较，YOLOv5-cls分类模型取得了不错的结果。

2. 背景

图像分类任务是计算机视觉领域的核心任务之一，其目标就是根据图像信息中所反映的不同特征，将不同类别的图像区分开来。钢厂在生产钢板的时候，由于工艺或者现场因素原因，有的钢板表面会产生缺陷，通过对钢板表面缺陷类别进行分类和统计，从而分析缺陷产生的原因，对进一步改善工艺，降低次品率有很大的帮助。

使用人工对每天生产的钢板进行缺陷检测和分类，不仅费时费力，而且很容易漏检和错检。基于计算机视觉的方法对图像进行分类现在已经很成熟了，目前比较主流的图像分类网络有ResNet、DenseNet、EfficientNet等。YOLOv5是目标检测方向的一个主流框架，但yolov5-6.2版本增加了图像分类功能，从官方公布的实验结果来看，取得了不错的效果，本项目首次基于yolov5训练钢板表面缺陷分类模型。

3. 数据

国内外的工业界和学术界目前开源了几个钢板表面缺陷数据集，说实话，部分数据集的中图片的数量和质量还是有待提高的，但是想一想整理并开源一个数据集耗费大量的人力和物力，成本是巨大的，而且这些数据都是商业数据，也是保密的，能够免费使用目前开源的数据已经很幸运了，且用且珍惜吧。其实本人也参与过钢板表面缺陷检测项目，深入钢厂在钢板生产现场收集了某个钢种的一些钢板表面缺陷图片，也和现场业务员进行沟通，对缺陷图片进行了标注，整理了一份包含1600张缺陷图片的数据集，总共4类缺陷，每类缺陷400张图片，但是这是商业数据，也是保密数据，无法公开。其实，我很想继续把这个数据集进行扩充，收集更多的缺陷图片，增加更多的缺陷类别，当有一天这个数据成果公开后，可以给工业界和学术界带来更多的图像分类、目标检测和图像分割研究成果，希望这一天早日到来吧。

本项目中使用的图像分类数据集，是东北大学带钢表面缺陷检测数据集，这个数据集在之前的“【表面缺陷检测】表面缺陷检测数据集汇总”介绍过，数据集中包含夹杂、划痕、压入氧化皮、裂纹、麻点和斑块6种缺陷，每种缺陷300张，图像尺寸为200×200。

英文名称	中文名称	图片数	图片尺寸
crazing	裂纹	300	200×200
inclusion	夹杂	300	200×200
patches	斑块	300	200×200
pitted_surface	麻点	300	200×200
rolled-in_scale	压入氧化皮	300	200×200
scratches	划痕	300	200×200

4. 代码

4.1 项目搭建

下载代码，将代码上传到服务器，也可以是使用本地的Windows系统进行训练：

我们只要用到的是classify文件夹下的train.py、val.py和predict.py脚本代码：

train.py：训练脚本
val.py：评估脚本
predict.py：推理脚本

4.2 核心代码

主干网络使用yolov5s、yolov5x、yolov5m、yolov5n、yolov5l：

class ClassificationModel(BaseModel):
    # YOLOv5 classification model
    def __init__(self, cfg=None, model=None, nc=1000, cutoff=10):  # yaml, model, number of classes, cutoff index
        super().__init__()
        self._from_detection_model(model, nc, cutoff) if model is not None else self._from_yaml(cfg)

    def _from_detection_model(self, model, nc=1000, cutoff=10):
        # Create a YOLOv5 classification model from a YOLOv5 detection model
        if isinstance(model, DetectMultiBackend):
            model = model.model  # unwrap DetectMultiBackend
        model.model = model.model[:cutoff]  # backbone
        m = model.model[-1]  # last layer
        ch = m.conv.in_channels if hasattr(m, 'conv') else m.cv1.conv.in_channels  # ch into module
        c = Classify(ch, nc)  # Classify()
        c.i, c.f, c.type = m.i, m.f, 'models.common.Classify'  # index, from, type
        model.model[-1] = c  # replace
        self.model = model.model
        self.stride = model.stride
        self.save = []
        self.nc = nc

    def _from_yaml(self, cfg):
        # Create a YOLOv5 classification model from a *.yaml file
        self.model = None

class Classify(nn.Module):
    # YOLOv5 classification head, i.e. x(b,c1,20,20) to x(b,c2)
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1):  # ch_in, ch_out, kernel, stride, padding, groups
        super().__init__()
        c_ = 1280  # efficientnet_b0 size
        self.conv = Conv(c1, c_, k, s, autopad(k, p), g)
        self.pool = nn.AdaptiveAvgPool2d(1)  # to x(b,c_,1,1)
        self.drop = nn.Dropout(p=0.0, inplace=True)
        self.linear = nn.Linear(c_, c2)  # to x(b,c2)

    def forward(self, x):
        if isinstance(x, list):
            x = torch.cat(x, 1)
        return self.linear(self.drop(self.pool(self.conv(x)).flatten(1)))

损失函数：

def smartCrossEntropyLoss(label_smoothing=0.0):
    # Returns nn.CrossEntropyLoss with label smoothing enabled for torch>=1.10.0
    if check_version(torch.__version__, '1.10.0'):
        return nn.CrossEntropyLoss(label_smoothing=label_smoothing)
    if label_smoothing > 0:
        LOGGER.warning(f'WARNING ⚠️ label smoothing {label_smoothing} requires torch>=1.10.0')
    return nn.CrossEntropyLoss()

优化器：

def smart_optimizer(model, name='Adam', lr=0.001, momentum=0.9, decay=1e-5):
    # YOLOv5 3-param group optimizer: 0) weights with decay, 1) weights no decay, 2) biases no decay
    g = [], [], []  # optimizer parameter groups
    bn = tuple(v for k, v in nn.__dict__.items() if 'Norm' in k)  # normalization layers, i.e. BatchNorm2d()
    for v in model.modules():
        for p_name, p in v.named_parameters(recurse=0):
            if p_name == 'bias':  # bias (no decay)
                g[2].append(p)
            elif p_name == 'weight' and isinstance(v, bn):  # weight (no decay)
                g[1].append(p)
            else:
                g[0].append(p)  # weight (with decay)

    if name == 'Adam':
        optimizer = torch.optim.Adam(g[2], lr=lr, betas=(momentum, 0.999))  # adjust beta1 to momentum
    elif name == 'AdamW':
        optimizer = torch.optim.AdamW(g[2], lr=lr, betas=(momentum, 0.999), weight_decay=0.0)
    elif name == 'RMSProp':
        optimizer = torch.optim.RMSprop(g[2], lr=lr, momentum=momentum)
    elif name == 'SGD':
        optimizer = torch.optim.SGD(g[2], lr=lr, momentum=momentum, nesterov=True)
    else:
        raise NotImplementedError(f'Optimizer {name} not implemented.')

    optimizer.add_param_group({'params': g[0], 'weight_decay': decay})  # add g0 with weight_decay
    optimizer.add_param_group({'params': g[1], 'weight_decay': 0.0})  # add g1 (BatchNorm2d weights)
    LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__}(lr={lr}) with parameter groups "
                f"{len(g[1])} weight(decay=0.0), {len(g[0])} weight(decay={decay}), {len(g[2])} bias")
    return optimizer

5. 训练

训练命令：

python classify/train.py --model weights/yolov5s-cls.pt --data data_custom --epochs 100  --batch-size 32 --imgsz 224

开始训练：

结束训练：

生成的文件：

我使用了一张Tesla P100显卡，训练了100 epoch，仅用了不到6分钟的时间，真的是相当的快，而且top1和top5精度都达到了100%。

6. 评估和推理

6.1 评估

评估代码：

python classify/val.py --weights runs/train-cls/exp/weights/best.pt --data data_custom

评估结果：

6.2 推理

推理代码：

# 测试im1.jpg
python classify/predict.py --weights runs/train-cls/exp/weights/best.pt --source im1.jpg

# 测试im2.jpg
python classify/predict.py --weights runs/train-cls/exp/weights/best.pt --source im2.jpg

推理结果：

7. 导出

7.1 ONNX

执行命令导出onnx：

python export.py --weights runs/train-cls/exp/weights/best.pt --include onnx

输出：

Detect:          python classify/predict.py --weights runs/train-cls/exp/weights/best.onnx
Validate:        python classify/val.py --weights runs/train-cls/exp/weights/best.onnx
PyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', 'runs/train-cls/exp/weights/best.onnx')  # WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference
Visualize:       https://netron.app

我们使用onnx进行模型部署，也可以把它当作中间件进行模型转换，根据项目需求进行选择。

7.2 TensorRT

执行命令导出engine：

python export.py --weights runs/train-cls/exp/weights/best.pt --include engine --device 0

输出：

Detect:          python classify/predict.py --weights runs/train-cls/exp/weights/best.engine
Validate:        python classify/val.py --weights runs/train-cls/exp/weights/best.engine
PyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', 'runs/train-cls/exp/weights/best.engine')  # WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference
Visualize:       https://netron.app

我们导出的engine模型，可以使用英伟达的TensorRT框架进行部署，加速模型推理速度。

8. 结论

本次项目基于yolov5对钢板表面缺陷进行分类，从评估指标来看，验证集上准确率很高，取得了很好的结果。模型训练速度也很快，模型导出也很方便和友好。目前，还可以做的工作是使用多GPU对模型进行训练，对导出的模型进行部署。但是，虽然分类精度很高，但是我们不知道缺陷的具体位置，没办法对缺陷进行定位，目前的想法是使用卷积热力图，用于突出图像的类的特定区域，不知道是否可行，还需进一步验证。总后絮叨，yolov5你真强！截止发文前，ultralytics公司目前已将yolov8开源了，yolov8将在江湖上引起腥风血雨。

如果您觉得这篇文章对您有一点点的帮助和启发，希望您关注公众号AIexplore，并点赞、转发。您可以联系我获取项目中的数据集和代码，数据整理和代码调试不易，公众号运营困难，为了公众号的正常运营，提供有偿指导，感谢您的理解和支持，祝好。

联系方式：公众号底部菜单栏--关于我--与我联系【订阅CSDN专栏的朋友，请加我v，发您数据和代码，不贴出代码和数据集的链接是为了防止爬虫，望理解】

参考资料

[1]https://github.com/ultralytics/yolov5
[2]https://blog.csdn.net/AugustMe/article/details/128111977

【图像分类】基于yolov5的钢板表面缺陷分类(附代码和数据集)

1. 引言

2. 背景

3. 数据

4. 代码

4.1 项目搭建

4.2 核心代码

5. 训练

6. 评估和推理

6.1 评估

6.2 推理

7. 导出

7.1 ONNX

7.2 TensorRT

8. 结论

参考资料

热门文章

最新文章

相关课程

相关电子书

相关实验场景

热门

活动广场

任务中心

开发者评测

高校计划

乘风者计划

训练营

阿里云MVP

话题

直播

下载

镜像站

技术资料

插件

【图像分类】基于yolov5的钢板表面缺陷分类(附代码和数据集)

1. 引言

2. 背景

3. 数据

4. 代码

4.1 项目搭建

4.2 核心代码

5. 训练

6. 评估和推理

6.1 评估

6.2 推理

7. 导出

7.1 ONNX

7.2 TensorRT

8. 结论

参考资料

热门文章

最新文章

相关课程

相关电子书

相关实验场景