Yolov5 钢材表面缺陷识别算法及代码实践

视觉/图片重量级资讯，第一时间送达！

来源 | 智能制造情报局

介绍

钢材缺陷包括表面质量、内部缺陷、化学成分偏析、力学性能不合格等。作为产品标准的要求，对钢材的表面质量有明确的规定，即钢材表面不得有裂纹、气泡、疤痕、夹杂、褶皱等。本文介绍了一种识别钢材表面缺陷的算法，并附有代码练习。

1）背景

我在东京大学读硕士的时候接触过NEU-DET数据集，这个数据集是用来识别钢铁表面缺陷的，最近想起来了，就试着用它来识别钢铁表面的缺陷。

数据集：

2）检测算法

关于算法的文章很多，我就不一一赘述了，下面是最近看到的几篇不错的文章：

YOLO系列（从v1到v5）模型解读（上、中、下）：

&&核心基础知识完整讲解（第1部分和第2部分）：

3）数据理解

东大钢铁缺陷数据汇集了热轧带钢的六种典型表面缺陷：

规模（RS）

（帕）

（铬）

（附注）

（在）

（雜誌）

可视化一些数据：

可以看出，该数据集的类内缺陷在外观上差异很大。

例如划痕（Sc）可能是水平划痕、垂直划痕、倾斜划痕等。另外，由于光照、材料变化等影响，类内缺陷图像的灰度是变化的。

简而言之，NEU数据集包含两个困难的挑战，即类内缺陷在外观上具有很大差异，而类间缺陷具有相似的方面。

4）数据处理

该数据集的标签是XML格式，需要转化为TXT格式，并且对数据进行归一化，才能供YOLO算法使用。

转换主要涉及校准框的规范化和类型的规范化。

import xml. etree. ElementTree as ETimport pickleimport osfrom os import listdir, getcwdfrom os. path import joinimport globclasses =["crazing", "inclusion", "patches", "pitted_surface","rolled-in_scale", "scratches"]
def convert(size, box):    dw=1./size[0]    dh =1./size[1]    x = (box[0] + box[1]) / 2.0    y=(box[2]+box[3])/2.0    w=box[1] -box[0]    h = box[3] - box[2]    x=x*dw    w=w*dw    y= y*dh    h =h*dh    return(x,y,w,h)
def convert_annotation(image_name):    in_file=open('./aragnal/ANNOTATIONS/'+image_name[: -3]+'xmL')    out_file=open( "D:/python/kaggle/yolov5-master/NEU-DET/valid/labels/"+image_name[: -3]+'txt', 'w')    tree=ET. parse(in_file)    root=tree. getroot()    size=root.find('size')    w=int(size.find('width').text)    h=int(size. find( 'height'). text)    for obj in root.iter('object'):        cls=obj. find('name').text        if cls not in classes:            print(cls)            continue        cls_id = classes.index(cls)        xmlbox = obj.find('bndbox')        b = (float(xmlbox.find('xmin').text),            float(xmlbox.find('xmax').text),            float(xmlbox.find('xmin').text),            float(xmlbox.find('xmax').text))        bb = convert((w, h), b)        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb])+'\n')
wd = getcwd()
if __name__ == '__main__':    for image_path in glob.glob(r"D:\python\kaggle\yolov5-master\NEU-DET\valid\images\*.jpg"):        image_name = image_path.split('\\')[-1]    # print (imege path)        convert_annotation(image_name)

结果如下：

5）调整参数

train.py配置参数如下：

--批量大小

-时代

300

- 数据

D:\\-\NEU-DET\data.yaml

–cfg

目录:\\-\-\\.yaml

–

D:\\-\-\runs\train\exp7\\best.pt

路径可以是绝对路径，也可以是相对路径，这里使用的是绝对路径。

batch-size和显卡显存有关，如果超出显存会报错，本次设置的是8。

纪元设定为300次。

6）结果预测

运行train.py训练集结果如下：

PR结果很差，一方面是因为机器不够好，epoch太少，导致训练不充分，欠拟合钢材的表面缺陷，另一方面可能和数据集中类间相似度低有关，COCO训练集的数据增强手段可能不完全适合这个项目，需要额外的数据增强手段。

.py预测参数如下：

–

D:\\-\NEU-DET\train\

–

D:\\-\-\runs\train\exp8\\best.pt

--conf

0.3

结果如下：

7）经验总结

1、如果第一次训练的结果比人眼的效果差很多，那么最有可能的原因是数据图像预处理出现了问题。

2. 有多个模型（,,,），对应衣服尺码smlx，不要只选最大的，要根据自己穿的衣服尺码来选（开个玩笑），实际原因是盲目选择过大的尺寸容易过拟合，尤其是这种数据量不大的项目。

3.由于数据集较小，我正在考虑进行迁移学习钢材的表面缺陷，但不太清楚如何冻结一些权重。

8）改进

该数据集的难度主要体现在类间差距不明显，类内差距较大，对上面的mAP 0.45进行了如下改进，最终mAP提升到0.61：

（最近有点忙，以后开个博客详细写）

1、由于钢材较细长，有些裂纹集中在一个方向，而划痕则是四面八方的。两种类型类似，所以特定方向的划痕很容易被误认为是裂纹。因此增加了裂纹的数据增强（自旋转拼接等）

2.由于图像主要呈现的是RGB的g，所以另外两个增强所增加的泛化能力没有意义，所以删除这个增强（同样有增加计算速度的效果）

3. 类别间差距过大，需要单独考虑数据增强，比如划痕在每幅图像中往往会出现好几次，可以考虑对划痕多做些拼接操作。

4. 在做拼接操作的时候，不能把数据想象成一个人，比如对于斑块这个训练类别，会把斑块拼接起来，形成一个新的车牌，这样如果单独根据两个斑块做一个数据框就会不准确（下图如果做上下拼接就是这种情况）