:568張圖如何用YOLO跑通產(chǎn)線)
簡(jiǎn)介面向工業(yè)視覺(jué)缺陷檢測(cè)場(chǎng)景的YOLO軸承生產(chǎn)缺陷數(shù)據(jù)集與配套腳本適合目標(biāo)檢測(cè)學(xué)習(xí)者、工業(yè)質(zhì)檢開(kāi)發(fā)者及希望復(fù)現(xiàn)YOLO訓(xùn)練流程的Python用戶(hù)使用。數(shù)據(jù)集包含568張軸承圖片標(biāo)注三類(lèi)常見(jiàn)生產(chǎn)缺陷預(yù)處理與標(biāo)注思路可直接借鑒配套txt標(biāo)簽與XML標(biāo)注可適配YOLO官方Darknet及YOLOv5等常見(jiàn)框架。壓縮包共1772個(gè)文件涵蓋624張JPG圖片、577個(gè)txt標(biāo)簽、568個(gè)XML標(biāo)注、2個(gè)Python腳本和1個(gè)YAML配置整體約755MB目錄結(jié)構(gòu)清晰可按圖像、標(biāo)簽、配置分層使用。已有337人學(xué)習(xí)下載。利用該資源可完整體驗(yàn)從數(shù)據(jù)準(zhǔn)備、標(biāo)注格式轉(zhuǎn)換、模型配置到缺陷檢測(cè)的訓(xùn)練鏈路Python腳本和YAML文件有助于快速啟動(dòng)訓(xùn)練并降低調(diào)參門(mén)檻同時(shí)也可用于研究軸承檢測(cè)中的光照變化、拍攝角度差異和背景干擾等魯棒性問(wèn)題。這套數(shù)據(jù)與代碼組合可作為軸承外觀質(zhì)檢和YOLO目標(biāo)檢測(cè)入門(mén)項(xiàng)目的起步材料。1. 568張圖訓(xùn)YOLO做軸承缺陷檢測(cè)到底靠不靠譜看到568張訓(xùn)練圖不少人的第一反應(yīng)是“數(shù)據(jù)量太少”。但在真實(shí)軸承產(chǎn)線上能整理出568張帶標(biāo)注的缺陷圖已經(jīng)是能啟動(dòng)項(xiàng)目的底子。軸承是標(biāo)準(zhǔn)件缺陷形態(tài)高度重復(fù)一張劃傷樣本能代表一整批同類(lèi)故障可也正因?yàn)槭菢?biāo)準(zhǔn)件光照、角度、反光變化會(huì)讓同一個(gè)缺陷長(zhǎng)得不一樣。這個(gè)標(biāo)題把YOLO、568張、三類(lèi)缺陷放在一起本質(zhì)是一個(gè)工業(yè)小樣本目標(biāo)檢測(cè)命題數(shù)據(jù)量有限、類(lèi)別明確、要求盡快跑通并讓產(chǎn)線接受。這篇筆記按數(shù)據(jù)準(zhǔn)備、模型選型、訓(xùn)練參數(shù)、避坑、部署驗(yàn)證的順序把這套方案拆完整。2. 把568張圖準(zhǔn)備成能訓(xùn)練的數(shù)據(jù)集三類(lèi)缺陷標(biāo)注與劃分策略2.1 三類(lèi)缺陷怎么定劃傷、麻點(diǎn)、剝落的標(biāo)注邊界軸承表面缺陷最常被歸成三類(lèi)劃傷scratch是細(xì)長(zhǎng)溝狀痕跡經(jīng)常與金屬反光混在一起麻點(diǎn)pitting是成片出現(xiàn)的細(xì)小凹坑單看輪廓不大但密度高剝落spalling是表面材料成塊掉落后留下的不規(guī)則區(qū)域邊界通常很毛糙。這三類(lèi)在可見(jiàn)光成像下特征差異明顯用YOLO做目標(biāo)檢測(cè)完全能分開(kāi)前提是標(biāo)注時(shí)不能抱著“差不多就行”的心態(tài)。標(biāo)注邊界是第一批坑。麻點(diǎn)通常成群出現(xiàn)逐個(gè)點(diǎn)去標(biāo)的話(huà)一張圖能標(biāo)出幾百個(gè)框標(biāo)到懷疑人生而且YOLO的NMS會(huì)把密集框互相抑制導(dǎo)致漏檢。常見(jiàn)做法是把一團(tuán)密集麻點(diǎn)當(dāng)作整體標(biāo)成一個(gè)框框住聚集區(qū)域劃傷則必須標(biāo)完整條軌跡不能只標(biāo)一小段剝落用外接矩形框住剝落塊框里帶進(jìn)一點(diǎn)正常表面也能接受。如果一張圖同時(shí)出現(xiàn)兩類(lèi)缺陷就分多個(gè)框標(biāo)每個(gè)框只給一個(gè)類(lèi)別。如果你用的是YOLO格式一張圖對(duì)應(yīng)一個(gè)同名txt文件每行是“class_id cx cy w h”坐標(biāo)全部歸一化到0到1。標(biāo)注工具用labelImg這類(lèi)能直接導(dǎo)出YOLO格式的即可labelme也行但輸出JSON還要轉(zhuǎn)一道。建議從第一天就直接存YOLO格式后面訓(xùn)練能省掉轉(zhuǎn)換腳本的麻煩。2.2 按工件分組劃分?jǐn)?shù)據(jù)集而不是按圖片隨機(jī)劃分568張圖看著少真正的坑往往不在數(shù)量而在劃分方式。假設(shè)一條產(chǎn)線上同一個(gè)軸承被拍了八張不同角度的圖把這八張圖隨機(jī)分到train和val里模型訓(xùn)練時(shí)已經(jīng)見(jiàn)過(guò)同一個(gè)表面的不同角度版本驗(yàn)證分?jǐn)?shù)會(huì)虛高到產(chǎn)線換新批次就崩。軸承這類(lèi)標(biāo)準(zhǔn)件正確的劃分維度是“工件”而不是“圖片”。我一般用按工件分組的分層采樣先把同一個(gè)軸承ID下所有圖片編號(hào)設(shè)成同一組再按每個(gè)類(lèi)別的占比做分層劃分。sklearn里可以這樣實(shí)現(xiàn)import numpy as np from sklearn.model_selection import StratifiedGroupKFold image_files [...] # 568張圖的路徑列表 group_ids [...] # 同一個(gè)軸承的多張圖填同一個(gè)ID例如 B001、B001、B002... labels [...] # 每張圖的主缺陷類(lèi)別 0/1/2 sgkf StratifiedGroupKFold(n_splits5, shuffleTrue, random_state42) for train_idx, val_idx in sgkf.split(image_files, labels, groupsgroup_ids): break # 取第一折相當(dāng)于約 8:2 的劃分 print(len(train_idx), len(val_idx))這段代碼的核心是StratifiedGroupKFold同時(shí)做了兩件事分層保證每一折里三類(lèi)樣本比例接近全量數(shù)據(jù)分組保證同一個(gè)軸承ID的圖片不會(huì)同時(shí)出現(xiàn)在train和val。如果一張圖包含多類(lèi)缺陷labels就填數(shù)量最多的那個(gè)類(lèi)別分組信息不受影響。劃分完記得把結(jié)果落成目錄常見(jiàn)做法是生成images/train和images/val兩個(gè)文件夾再把對(duì)應(yīng)txt按同樣路徑復(fù)制過(guò)去。驗(yàn)證集給到總數(shù)據(jù)的兩成也就是110張左右如果樣本實(shí)在太少也可以從train里再切一部分當(dāng)test測(cè)試集只在最終評(píng)估時(shí)用一次訓(xùn)練過(guò)程中反復(fù)看測(cè)試集是給自己留雷。2.3 小樣本增強(qiáng)策略在線增強(qiáng)要保守離線增強(qiáng)要克制YOLO默認(rèn)開(kāi)啟的在線增強(qiáng)包括mosaic、mixup、HSV擾動(dòng)、隨機(jī)翻轉(zhuǎn)和縮放這些是針對(duì)大數(shù)據(jù)集設(shè)計(jì)的。568張做基數(shù)時(shí)mosaic把四張圖拼一起缺陷目標(biāo)在拼接后被縮小甚至裁掉一半對(duì)小目標(biāo)檢測(cè)反而不利。我的習(xí)慣是mosaic降到0.3左右mixup降到0.1或者直接關(guān)掉上下翻轉(zhuǎn)要看軸承擺放方向產(chǎn)線相機(jī)角度固定就關(guān)掉保留左右翻轉(zhuǎn)、輕微旋轉(zhuǎn)和亮度對(duì)比度擾動(dòng)。離線增強(qiáng)可以補(bǔ)量但別無(wú)腦擴(kuò)。旋轉(zhuǎn)角度控制在5度到15度再大就違背軸承的實(shí)際成像規(guī)律高斯噪聲、亮度變化、對(duì)比度拉伸這些模擬光照波動(dòng)的增強(qiáng)性?xún)r(jià)比最高。離線增強(qiáng)要在數(shù)據(jù)集劃分完成之后單獨(dú)對(duì)train做驗(yàn)證集和測(cè)試集必須保持原始圖像否則增強(qiáng)過(guò)的紋理會(huì)被模型當(dāng)成特征驗(yàn)證分?jǐn)?shù)虛高。見(jiàn)過(guò)有人把訓(xùn)練集擴(kuò)到兩三千張mAP看著不錯(cuò)一下產(chǎn)線就翻車(chē)最后查出來(lái)驗(yàn)證集也被增強(qiáng)過(guò)評(píng)估結(jié)果根本不代表真實(shí)水平。注意產(chǎn)線上除了缺陷圖還會(huì)出現(xiàn)大量正常件建議單獨(dú)留一批OK工件圖做誤檢驗(yàn)證不要混進(jìn)缺陷訓(xùn)練集。訓(xùn)練時(shí)沒(méi)有負(fù)樣本不代表部署時(shí)沒(méi)有誤檢率要用真實(shí)產(chǎn)線數(shù)據(jù)單獨(dú)統(tǒng)計(jì)。3. YOLO模型選型與訓(xùn)練參數(shù)小數(shù)據(jù)集下從n到s的務(wù)實(shí)配置3.1 為什么小數(shù)據(jù)集先選YOLOv8n或YOLOv5n而不是大模型標(biāo)題只寫(xiě)了“基于YOLO”沒(méi)限定版本實(shí)際項(xiàng)目里用得最多的還是YOLOv8和YOLOv5兩個(gè)系列。對(duì)568張這個(gè)量級(jí)模型容量的選擇比版本選擇重要得多。YOLOv8n參數(shù)量在3M級(jí)別YOLOv8s在11M級(jí)別再往上在這個(gè)數(shù)據(jù)量下基本必過(guò)擬合。小模型收斂快、顯存占用低、部署也輕松而軸承缺陷模式高度重復(fù)n和s的精度差距沒(méi)有想象中那么大。我習(xí)慣先用YOLOv8n跑通完整流程把數(shù)據(jù)劃分、標(biāo)注質(zhì)量、評(píng)估指標(biāo)都理順再換YOLOv8s對(duì)比一次。如果兩類(lèi)提升明顯就留下來(lái)如果只漲零點(diǎn)幾個(gè)點(diǎn)果斷用n后面做TensorRT部署時(shí)幀率差距非常直觀。YOLOv5n也完全可以如果你有一張老卡要跑推理v5的部署生態(tài)更成熟。小模型另一個(gè)好處是對(duì)增強(qiáng)關(guān)閉更寬容。大模型需要更強(qiáng)增強(qiáng)壓過(guò)擬合小模型容量有限增強(qiáng)稍微保守反而收斂更干凈。3.2 用YOLOv8訓(xùn)練自己的軸承數(shù)據(jù)集最小可復(fù)現(xiàn)命令動(dòng)手訓(xùn)練前先把數(shù)據(jù)配置文件寫(xiě)好。data.yaml放在項(xiàng)目根目錄path: /work/bearing train: images/train val: images/val names: 0: scratch 1: pitting 2: spallingnc可以不寫(xiě)YOLO會(huì)從names長(zhǎng)度推斷類(lèi)別數(shù)。類(lèi)別名用英文是為了避免輸出層編碼問(wèn)題也方便后面做TensorRT導(dǎo)出時(shí)映射類(lèi)別。path建議用絕對(duì)路徑訓(xùn)練時(shí)能少踩一個(gè)“相對(duì)路徑找不到圖片”的坑。訓(xùn)練命令如下yolo detect train \ data/work/bearing/data.yaml \ modelyolov8n.pt \ epochs150 \ imgsz640 \ batch16 \ patience30 \ optimizerAdamW \ lr00.001 \ ampTrue \ project/work/bearing/runs參數(shù)按順序說(shuō)明data指向剛才的yamlmodel用yolov8n.pt預(yù)訓(xùn)練權(quán)重別從零開(kāi)始訓(xùn)epochs給150小數(shù)據(jù)集在100到200輪之間就能收斂再長(zhǎng)就是背訓(xùn)練集噪聲imgsz先固定640后面因?yàn)樾∧繕?biāo)漏檢再考慮提960batch在T4這種16G顯存上用AMP能開(kāi)到16甚至32顯存小就降到8patience30表示驗(yàn)證指標(biāo)連續(xù)30輪不提升就早停優(yōu)化器我用AdamW而不是默認(rèn)SGD小數(shù)據(jù)集下AdamW收斂更平滑lr0給0.001偏保守夠用。訓(xùn)練完看runs/detect/train/weights下的best.pt和last.ptbest是按驗(yàn)證集指標(biāo)保存的部署用best。第一次訓(xùn)練不要急著調(diào)參先看清日志里的P、R、mAP50三個(gè)數(shù)再?zèng)Q定動(dòng)哪里。3.3 損失函數(shù)與關(guān)鍵超參哪些值得動(dòng)哪些別亂動(dòng)YOLOv8默認(rèn)邊界框損失是CIoU分類(lèi)損失是BCE這對(duì)多數(shù)缺陷檢測(cè)夠用。很多教程讓換SIoU或者加focal loss但在568張規(guī)模下?lián)Q損失函數(shù)帶來(lái)的提升通常不如修正標(biāo)注和增強(qiáng)策略。我傾向于把損失函數(shù)當(dāng)最后手段不是第一優(yōu)化項(xiàng)。需要?jiǎng)拥氖侨齻€(gè)地方。第一是imgsz麻點(diǎn)這類(lèi)小目標(biāo)漏檢多時(shí)從640提到960mAP往往直接漲三五個(gè)點(diǎn)代價(jià)是顯存和推理時(shí)間上升。第二是類(lèi)別樣本量失衡比如劃傷300張、麻點(diǎn)只有80張YOLOv8命令行沒(méi)直接暴露類(lèi)別權(quán)重參數(shù)常見(jiàn)做法是用Python API自定義loss或者對(duì)少數(shù)類(lèi)做離線過(guò)采樣。第三是mosaic前文說(shuō)過(guò)小數(shù)據(jù)集、小目標(biāo)場(chǎng)景下mosaic經(jīng)常起反作用直接降權(quán)重最省事。還有一個(gè)容易忽略的點(diǎn)batch太小會(huì)讓BatchNorm統(tǒng)計(jì)不穩(wěn)顯存不夠時(shí)不要硬開(kāi)batch4不如用accumulate等效增大batch或者減小輸入分辨率保證batch能上去。訓(xùn)練時(shí)看一眼loss曲線的震蕩幅度如果劇烈跳動(dòng)先懷疑batch過(guò)小或lr過(guò)大這兩個(gè)參數(shù)的優(yōu)先級(jí)高于換損失函數(shù)。4. 避坑小數(shù)據(jù)集訓(xùn)YOLO的五個(gè)翻車(chē)點(diǎn)與排查方法4.1 訓(xùn)練loss一直在降驗(yàn)證mAP卻紋絲不動(dòng)現(xiàn)象日志里box_loss和cls_loss一路向下表面看訓(xùn)練很健康但每個(gè)epoch結(jié)束后的mAP50卡在0.2甚至更低。原因八成是增強(qiáng)過(guò)強(qiáng)模型學(xué)到了增強(qiáng)引入的偽紋理也可能標(biāo)注框偏移大、類(lèi)別標(biāo)錯(cuò)loss主要在擬合噪聲。兩種情況的共同特征是訓(xùn)練loss降得快驗(yàn)證指標(biāo)不動(dòng)。解決先把mosaic降到0、mixup關(guān)掉、亮度擾動(dòng)減弱用最樸素的數(shù)據(jù)跑一輪再對(duì)比train loss和val loss的gapgap越拉越大就是過(guò)擬合砍增強(qiáng)或換更小的模型提前早停。如果是標(biāo)注噪聲問(wèn)題抽20張圖人工復(fù)核標(biāo)注框這條檢查比調(diào)任何參數(shù)都有效。4.2 細(xì)小缺陷漏檢嚴(yán)重召回率上不去現(xiàn)象劃傷、麻點(diǎn)這類(lèi)小目標(biāo)在測(cè)試集上大量漏檢輸出的框位置基本對(duì)但recall只有0.3到0.4。原因YOLO的下采樣倍數(shù)對(duì)極小目標(biāo)不友好640輸入下小目標(biāo)特征經(jīng)過(guò)幾輪下采樣后只剩幾個(gè)像素密集麻點(diǎn)互相靠近NMS還會(huì)把相鄰框合并成一個(gè)框。解決先試imgsz960或1280通常立竿見(jiàn)影顯存不夠就用切片推理把原圖切成帶重疊的塊分別檢測(cè)再合并也就是SAHI的常見(jiàn)思路。切片尺寸一般取640重疊率20%推理時(shí)間會(huì)增加但召回率提升明顯。機(jī)油蓋、齒輪這類(lèi)小金屬件的缺陷檢測(cè)也常遇到同一個(gè)問(wèn)題處理方式基本一致。4.3 驗(yàn)證集分?jǐn)?shù)比訓(xùn)練集還高先查數(shù)據(jù)泄漏現(xiàn)象訓(xùn)練集mAP50是0.7驗(yàn)證集反而0.85怎么看都不合常理。原因幾乎可以肯定是同一個(gè)工件ID的圖片同時(shí)進(jìn)了train和val模型記住了該工件表面的紋理驗(yàn)證時(shí)直接套答案。軸承是標(biāo)準(zhǔn)件同批次表面紋理高度一致這種泄漏很難靠肉眼發(fā)現(xiàn)。解決回到2.2用StratifiedGroupKFold按工件分組重新劃分再訓(xùn)練。這個(gè)坑查起來(lái)最費(fèi)時(shí)間因?yàn)閿?shù)據(jù)量少人工翻圖很難看出哪張圖兩邊都出現(xiàn)過(guò)不如一開(kāi)始就把劃分邏輯寫(xiě)對(duì)。4.4 三類(lèi)中有一類(lèi)AP幾乎為0現(xiàn)象劃傷和剝落的AP都到0.7以上麻點(diǎn)那一類(lèi)只有0.05。原因該類(lèi)別樣本量太少或者標(biāo)注方式有問(wèn)題。比如麻點(diǎn)被標(biāo)成一個(gè)巨大的框正樣本質(zhì)量低模型學(xué)不到該類(lèi)邊界特征。解決先檢查該類(lèi)圖片數(shù)量和標(biāo)注框尺寸分布。樣本少就做離線過(guò)采樣把包含該類(lèi)的小圖區(qū)域復(fù)制粘貼到空白圖上擴(kuò)到和最多類(lèi)一個(gè)量級(jí)框太大就重新標(biāo)把密集麻點(diǎn)拆成多個(gè)緊貼的小框。數(shù)據(jù)補(bǔ)齊后重訓(xùn)比換損失函數(shù)有效得多。負(fù)樣本也要注意訓(xùn)練集全是缺陷圖沒(méi)有OK件圖部署時(shí)正常件的紋理、油污會(huì)被誤判成缺陷所以單獨(dú)準(zhǔn)備一批OK圖做驗(yàn)證十分必要。4.5 TensorRT FP16導(dǎo)出后精度明顯掉點(diǎn)現(xiàn)象PyTorch里mAP50有0.85導(dǎo)出TensorRT FP16后直接掉到0.7有時(shí)還在小目標(biāo)上瘋狂誤檢。原因FP16的數(shù)值范圍和精度低于FP32小目標(biāo)的框坐標(biāo)回歸對(duì)這些數(shù)值誤差更敏感測(cè)試集本身有邊界case時(shí)FP16會(huì)放大這種不穩(wěn)定。解決先跑一遍FP32引擎確認(rèn)不是數(shù)據(jù)泄漏導(dǎo)致的虛高再看掉點(diǎn)幅度。掉點(diǎn)超過(guò)一個(gè)百分點(diǎn)就用INT8加校準(zhǔn)集INT8也穩(wěn)不住就把敏感層保留FP32。掉點(diǎn)在合理范圍內(nèi)且誤檢率可接受FP16的提速優(yōu)勢(shì)是值得換的但上線前必須用驗(yàn)證集完整比對(duì)一輪不能只看幾張圖的視覺(jué)效果。5. 上線前的最后一公里TensorRT加速與按類(lèi)統(tǒng)計(jì)漏檢率5.1 導(dǎo)出ONNX并生成TensorRT引擎訓(xùn)練完的best.pt是PyTorch權(quán)重產(chǎn)線推理一般不用PyTorch直接跑。先導(dǎo)出ONNX再轉(zhuǎn)TensorRTyolo export modelbest.pt formatonnx imgsz640 opset12 trtexec --onnxbest.onnx --saveEnginebest.engine --fp16trtexec執(zhí)行完會(huì)在終端打印引擎的mean inference time這是估算并發(fā)路數(shù)的第一手?jǐn)?shù)據(jù)。在T4這類(lèi)16G卡上常見(jiàn)結(jié)果在幾毫秒到十幾毫秒之間具體由模型大小和輸入分辨率決定拿自己訓(xùn)練的模型實(shí)測(cè)最準(zhǔn)。5.2 單卡能帶幾路視頻流用延遲倒推并發(fā)很多人搜“T4 1080P 25幀每秒用TensorRT YOLO 640分辨率能支持多少路”這個(gè)問(wèn)題沒(méi)有固定答案只能按自己的引擎實(shí)測(cè)倒推。假設(shè)RTSP拉流后每幀切到640×640送入引擎單幀耗時(shí)t毫秒單卡理論吞吐就是1000/t FPS。目標(biāo)要求每條流25FPS理論路數(shù)是(1000/t)除以25再按50%到70%的利用率扣除取幀、預(yù)處理、NMS和調(diào)度開(kāi)銷(xiāo)得到實(shí)際可帶路數(shù)。舉例測(cè)得t5ms理論200FPS除以25FPS等于8路按60%利用率算大概穩(wěn)帶5路。不要滿(mǎn)打滿(mǎn)算推理卡滿(mǎn)后CPU取幀和圖像編解碼會(huì)最先扛不住。多路并發(fā)時(shí)單獨(dú)起取幀線程不要讓推理線程等網(wǎng)絡(luò)幀。5.3 用三類(lèi)召回率評(píng)估缺陷檢測(cè)的實(shí)際價(jià)值最終評(píng)估不是只看整體mAP要按類(lèi)別分別看召回率。軸承缺陷檢測(cè)里漏檢比誤檢貴一個(gè)缺陷滑過(guò)去可能直接造成整批退貨。我習(xí)慣用下面這段腳本對(duì)測(cè)試集做最終評(píng)估from ultralytics import YOLO model YOLO(runs/detect/train/weights/best.pt) metrics model.val(datadata.yaml, splittest, conf0.25, iou0.45) print(各類(lèi)別 mAP50-95:, metrics.box.maps) print(整體 precision/recall:, metrics.box.p, metrics.box.r)conf降到0.1時(shí)召回率還能往上走但誤檢也會(huì)增多。上線前用一段時(shí)間真實(shí)產(chǎn)線的OK工件圖統(tǒng)計(jì)誤檢率只報(bào)單一閾值下的數(shù)字沒(méi)有參考意義。我現(xiàn)在的習(xí)慣是拿到這類(lèi)小樣本項(xiàng)目先把數(shù)據(jù)按工件分組切好再固定用輕量模型跑通baseline確認(rèn)增強(qiáng)策略不激進(jìn)最后才談TensorRT和并發(fā)路數(shù)。這個(gè)順序幫我少踩了很多坑希望幫到你。本文還有配套的精品資源點(diǎn)擊獲取