YOLOv5實現(xiàn)鋼材表面缺陷檢測

目前，基于機器視覺的表面缺陷已經(jīng)在各個工業(yè)領域廣泛取代人工視覺檢測，包括3C、汽車、家電、機械制造、半導體電子、化工、制藥、航空航天、輕工等行業(yè)。許多基于深度學習的缺陷檢測方法也被廣泛應用于各種工業(yè)場景。

本文的代碼實戰(zhàn)，是基于YOLOv5目標檢測算法，在NEU表面缺陷數(shù)據(jù)集上實現(xiàn)鋼材表面缺陷檢測。要求Python>=3.7.0，PyTorch>=1.7。

一、YOLOv5

選取YOLOv5，一方面是因為從最終效果來看YOLOv5已經(jīng)相當優(yōu)秀，是發(fā)展的比較完備、使用比較廣泛的一個版本；而更重要的是YOLOv5的調(diào)用、訓練和預測都十分方便，為初學者提供了良好的練手工具。YOLOv5的另一個特點就是它為不同的設備需求和不同的應用場景提供了大小和參數(shù)數(shù)量不同的網(wǎng)絡。

如圖所示，大一點的模型比如YOLOv5l和YOLOv5x參數(shù)更多，在大規(guī)模的COCO數(shù)據(jù)集上有更高的預測準確率；而小模型比如YOLOv5n或YOLOv5s占空間更小，適合部署在移動設備上，且推理速度更快，適合于高幀 率視頻的實時檢測。

二、鋼材表面缺陷數(shù)據(jù)集

東北大學(NEU)表面缺陷數(shù)據(jù)集，收集了熱軋帶鋼6種典型的表面缺陷，即軋內(nèi)垢(RS)、斑塊(Pa)、裂紋(Cr)、點蝕面(PS)、夾雜物(In)和劃痕(Sc)。該數(shù)據(jù)庫包括1800張灰度圖像：6種不同類型的典型表面缺陷各300個樣本。

下圖為6種典型表面缺陷的樣本圖像，每張圖像的原始分辨率為200×200像素。從圖中，我們可以清楚地觀察到類內(nèi)缺陷在外觀上存在較大差異，例如劃痕(最后一列)可能是水平劃痕、垂直劃痕和傾斜劃痕等。與此同時，類間缺陷也具有相似的特征，如滾積垢、裂紋和坑狀表面。

此外，由于光照和材料變化的影響，類內(nèi)缺陷圖像的灰度會發(fā)生變化?？傊?，NEU表面缺陷數(shù)據(jù)庫包含兩個難題，即類內(nèi)缺陷存在較大外觀差異，類間缺陷具有相似方面，缺陷圖像受到光照和材料變化的影響。

對于缺陷檢測任務，數(shù)據(jù)集提供了標注，標注了每個圖像中缺陷的類別和位置。下圖展示了數(shù)據(jù)集上的一些檢測結(jié)果示例。對于每個缺陷，黃色框是表示其位置的邊界框，綠色標簽是類分數(shù)。

三、YOLOv5數(shù)據(jù)格式

YOLOv5標簽文件中每一行的數(shù)據(jù)為class, x, y, w, h，其中class是該物體的類別，x,y是檢測框中心坐標，w,h是檢測框的寬和高。

四、格式轉(zhuǎn)換voc2yolo

從網(wǎng)上獲取一些目標檢測的數(shù)據(jù)集資源標簽的格式都是VOC(xml格式)的，而YOLOv5訓練所需要的文件格式是YOLO(txt格式)的，這里就需要對xml格式的標簽文件轉(zhuǎn)換為txt文件。

同時訓練自己的YOLOv5檢測模型的時候，數(shù)據(jù)集需要劃分為訓練集和驗證集。這里提供了一份代碼將xml格式的標注文件轉(zhuǎn)換為txt格式的標注文件，并按比例劃分為訓練集和驗證集。

importxml.etree.ElementTreeasET
importos
fromosimportgetcwd
fromtqdmimporttqdm

classes=["crazing","inclusion","patches","pitted_surface","rolled-in_scale","scratches"]

defconvert(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)

defconvert_annotation(image_id):
in_file='./xml/%s.xml'%(image_id)
out_file=open('./labels/%s.txt'%(image_id),'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)

forobjinroot.iter('object'):
difficult=obj.find('Difficult').text
cls=obj.find('name').text
ifclsnotinclassesorint(difficult)==1:
continue
cls_id=classes.index(cls)
xmlbox=obj.find('bndbox')
b=(float(xmlbox.find('xmin').text),float(xmlbox.find('xmax').text),float(xmlbox.find('ymin').text),
float(xmlbox.find('ymax').text))
bb=convert((w,h),b)
out_file.write(str(cls_id)+""+"".join([str(a)forainbb])+'
')

if__name__=="__main__":
wd=getcwd()
print(wd)
ifnotos.path.exists('./labels/'):
os.makedirs('./labels/')
image_ids=os.listdir('./datasets')
forimage_idintqdm(image_ids):
convert_annotation(image_id.split('.')[0])

五、訓練

YOLOv5是通過yaml格式的配置文件來找到對應的訓練測試數(shù)據(jù)，因此在訓練前需要對該文件進行簡單配置。

我們可以通過修改YOLOv5的配置文件來實現(xiàn)，需要修改的配置文件有兩個，我們先復制一份data/coco.yaml，這里將其重命名為my_coco.yaml

在download前加上一個#注釋掉這段代碼

將train、val、test修改為自己的路徑，以train為例，NEU-DET/train/images/

將nc修改為數(shù)據(jù)的類別數(shù)，如鋼材表面缺陷，故修改為6

將names修改自己數(shù)據(jù)的類名，如names: ["crazing", "inclusion", "patches", "pitted_surface", "rolled-in_scale", "scratches"]

下一個需要修改的文件為模型配置文件，在models文件夾中有四個模型的配置文件：yolov5s.yaml、yolov5m.yaml、yolov5l.yaml和yolov5x.yaml，可以根據(jù)需要選擇相應的模型，這里以yolo5s.yaml為例，打開文件，修改文件中的nc為自己的類別數(shù)即可。

pythontrain.py--weightsweights/yolov5s.pt--cfgmodels/yolov5s.yaml--datadata/my_coco.yaml--epoch30--batch-size32

六、測試

訓練結(jié)束后，在runs/exp0文件夾下會自動生成訓練結(jié)果，其中包括模型權(quán)重、混淆矩陣、PR曲線等。

進行測試時，打開detect.py文件，修改source為檢測圖片的文件夾，運行之后，在run->detect中輸出檢測結(jié)果；或者通過如下方式進行測試。

pythondetect.py--source./data/images/example.jpg--weightsruns/exp0/weights/best.pt--conf-thres0.25

最終預測的結(jié)果如下，可以看出模型缺陷檢測效果。

審核編輯：劉清