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

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

一、YOLOv5

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

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

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

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

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

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

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

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

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

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

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

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

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])

五、訓(xùn)練

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

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

在download前加上一個(gè)#注釋掉這段代碼

將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"]

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

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

六、測試

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

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

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

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

審核編輯：劉清