YOLOv5 實(shí)現(xiàn)無(wú)人機(jī)識(shí)別
點(diǎn)擊下方“AI算法與圖像處理”,一起進(jìn)步!
重磅干貨,第一時(shí)間送達(dá)
編輯: gloomyfish
作者: 陳金光
引言
YOLOv5發(fā)布至今,已經(jīng)迭代多個(gè)版本,是唯一基于pytorch框架的YOLO系列的模型框。簡(jiǎn)單易用,特別適合自定義目標(biāo)檢測(cè)與對(duì)象檢測(cè)。只要用了Pytorch版本YOLOv5框架,可以毫不夸張的說(shuō)兩個(gè)小時(shí)就可以學(xué)會(huì)目標(biāo)檢測(cè)。
01
第一步數(shù)據(jù)準(zhǔn)備
想用視覺(jué)識(shí)別一下空中飛行的無(wú)人機(jī),識(shí)別對(duì)象有:
1:鳥類bird2:無(wú)人機(jī)UAV3:直升機(jī)helicopter
想通過(guò)自己創(chuàng)建一個(gè)數(shù)據(jù)集,訓(xùn)練yolov5,在調(diào)用detect.py來(lái)識(shí)別一下效果。
百度了一下,發(fā)現(xiàn)無(wú)法獲取別人已經(jīng)標(biāo)注好的數(shù)據(jù)集。(這里提醒大家,自己做數(shù)據(jù)集不費(fèi)時(shí)間,給100多張圖打標(biāo)簽label就行,用不了半個(gè)小時(shí))圖片就從百度搜索就行,我們的目的是感性的認(rèn)識(shí)這套牛掰的算法,100張圖訓(xùn)練個(gè)幾個(gè)小時(shí)就可以出來(lái)很好的效果,還有不用非要顯卡,雖然我的電腦有顯卡,但是裝pytorch的顯卡cuda版本的總是出錯(cuò),索性不弄了,直接用cpu版本,cpu版本的pytorch包還小,一會(huì)兒就能安裝好,這100張圖我訓(xùn)練了100個(gè)epoch花費(fèi)了3到4個(gè)小時(shí),我的電腦的cpu是i7-9700的)
我是從百度上找了大概130多張圖片,我截個(gè)圖:

用labelimage打上標(biāo)簽,所有的標(biāo)簽數(shù)量大概也就200來(lái)個(gè)吧!不會(huì)用labelimgage制作數(shù)據(jù)集的同學(xué)可以看這里:
https://github.com/tzutalin/labelImgyolo的數(shù)據(jù)文件夾格式又要鞏固一下了,先建一個(gè)數(shù)據(jù)文件夾data,在數(shù)據(jù)文件夾data目錄下面建圖像文件夾images和標(biāo)簽文件夾labels,然后在images文件夾下面建立訓(xùn)練數(shù)據(jù)集train和驗(yàn)證集val(validation的縮寫,驗(yàn)證校驗(yàn)的意思),在labels的文件夾下面也是建立train和val兩個(gè)文件夾,分別放訓(xùn)練集和驗(yàn)證集的標(biāo)簽。
另外說(shuō)一下驗(yàn)證集,驗(yàn)證集里面我找了26張圖片,標(biāo)記了40個(gè)左右的labels吧,所以所有的數(shù)據(jù)加起來(lái)大概150張圖片。
我放一個(gè)圖片看一下驗(yàn)證集里面的圖片如下圖:

到這里為止,所有的數(shù)據(jù)基本整理完畢,另外還需要幾張測(cè)試的圖片,到測(cè)試的時(shí)候再?gòu)木W(wǎng)上找就是了。
02
配置與修改參數(shù)
這里下載yolov5的源碼就不寫怎么下載了,自己百度就好,相對(duì)難度低,我是從gitee上下載的,github總是卡,這個(gè)源碼也沒(méi)有多大,幾分鐘的時(shí)間就下載好了。下面直接寫怎么改一下代碼訓(xùn)練自己制作好的無(wú)人機(jī)的數(shù)據(jù)集。
分步驟寫吧
首先:明確思路
怎么訓(xùn)練這個(gè)算法?我要訓(xùn)練什么模型?我要拿這個(gè)模型做什么?這些問(wèn)題是在訓(xùn)練這個(gè)算法的時(shí)候首先要想的問(wèn)題。
回答:
我要訓(xùn)練一個(gè)模型能夠識(shí)別空中飛行的無(wú)人機(jī),并打上紅框label,空中的復(fù)雜環(huán)境下可能不只有無(wú)人機(jī),會(huì)有鳥、樹、白云、直升機(jī)、飛機(jī)等等其他干擾因素印象,簡(jiǎn)單起見,我又加上了鳥bird和直升機(jī)helicopter,所以用yolov5分類nc(number class)為3個(gè),也就是無(wú)人機(jī)UAV、鳥bird、直升機(jī)、helicopter,看下圖,建立一個(gè)數(shù)據(jù)信息文件,文件格式y(tǒng)aml

我把這個(gè)文件名保存成了chenchen.yaml
訓(xùn)練集圖片train在uav_train下面的data里面,訓(xùn)練集也是,這個(gè)文件不用寫標(biāo)簽label的地址,因?yàn)槲募A結(jié)構(gòu)是定的,程序會(huì)自動(dòng)判斷出label的位置。
chenchen.yaml文件里面的nc改成2,names改成自己的,也就是uav bird 和helicopter
下面我貼兩張圖片,這兩張圖片是我訓(xùn)練了100張圖后,測(cè)試的標(biāo)簽和測(cè)試的預(yù)測(cè),基本上框的位置已經(jīng)很好了,只有左上角的這個(gè)直升機(jī)helicopter沒(méi)有被識(shí)別出來(lái),這個(gè)我感覺(jué)是因?yàn)槲艺业闹鄙龣C(jī)圖片的問(wèn)題,這個(gè)直升機(jī)圖片有點(diǎn)特別,訓(xùn)練數(shù)據(jù)中這樣的圖片少的原因。

總結(jié):思路明確了,就是做一個(gè)三分類模型,無(wú)人機(jī)、鳥、直升機(jī),
其次修改train.py文件
思路確定了后就是訓(xùn)練這個(gè)模型了,我們已經(jīng)把數(shù)據(jù)集準(zhǔn)備好了,現(xiàn)在就訓(xùn)練這個(gè)模型。下載下來(lái)的yolo源碼里面的train.py文件,我截圖一下:

這個(gè)是train.py文件
第一:導(dǎo)入各種依賴包;第二:定義train函數(shù)方法,主要的就是這個(gè)train方法;第三:主程序;
解釋一下,各種依賴包不用解釋,train方法就是作者的核心了,簡(jiǎn)單解釋就是yolo算法的原理,主程序里面就是調(diào)用tran方法畫框框,(到目前位置我能給的解釋也只有這么多,我沒(méi)有花費(fèi)時(shí)間研究,我感性上的理解應(yīng)該是這樣)
這個(gè)文件中主要要修改的參數(shù)是什么呢?
下面這張圖中我給出了箭頭

第一個(gè)箭頭:
yolov5s.pt,這個(gè)是權(quán)重文件,用的是s版本,small,yolov5的權(quán)重文件不同分為幾個(gè)不同版本,這里我用的是s版本,這里可以認(rèn)為是模型的初始權(quán)重,
第二個(gè)箭頭:
這個(gè)yolov5s.yaml文件我保存到了uav_train文件夾里面,這個(gè)文件是配置yolov5模型的,叫做配置文件吧,等下這個(gè)文件我截圖一下,這個(gè)文件也要改一個(gè)地方。
第三個(gè)箭頭:
就是數(shù)據(jù)文件,就是上面定義的一個(gè)chenchen.yaml文件,這里有訓(xùn)練數(shù)據(jù)集的地址。
第四個(gè)箭頭:
是epoch,我寫的是100,大家也可以改成10試一下,20分鐘訓(xùn)練完成,我當(dāng)時(shí)做了一下,10個(gè)epoch效果很差,100個(gè)的時(shí)候效果就是我視頻的效果,都能到70%準(zhǔn)確率以上。
第五個(gè)箭頭:
是batch,我是16張圖片一個(gè)batch
第二個(gè)箭頭解釋圖片如下:

這里是這個(gè)算法的配置文件,最終輸出的類別是3,這改成3保存就行,其他的能看出來(lái)這個(gè)里面是yolo的架構(gòu),我也不會(huì)解釋,有想了解的可以自己百度。
總結(jié):到這里train.py就修改完畢了
03
訓(xùn)練與測(cè)試
到這里就可以訓(xùn)練了,我用的是pycharm這個(gè)集成環(huán)境,所以選中train.py文件直接右鍵運(yùn)行就可以

下面就是漫長(zhǎng)的訓(xùn)練時(shí)間,只有等著
其實(shí)這一步可以不用,test.py文件主要是用來(lái)測(cè)試驗(yàn)證集val的,這里和訓(xùn)練一樣,不要把訓(xùn)練好的權(quán)重文件的地址換一下,就是best.pt文件,data的參數(shù)改一下,改成chenchen.yaml就可以。最激動(dòng)人心的是第五步,就是訓(xùn)練好的模型驗(yàn)證一下其他的圖片,看看效果怎么樣。
測(cè)試新圖片,新視頻看效果
這個(gè)文件是detect.py,我截圖看下:

這里有四個(gè)箭頭
第一個(gè):訓(xùn)練好的最好的權(quán)重文件,加載上第二個(gè):我從百度找了一個(gè)視頻文件保存成了666.mp4第三個(gè):置信度,我這里其實(shí)可以寫0.7,第四個(gè):iou,
解釋:如果訓(xùn)練的epoch少,可以把三四箭頭改小點(diǎn),要不然視頻和圖片中不會(huì)被標(biāo)記出,會(huì)讓你誤認(rèn)為自己訓(xùn)練的不對(duì)。測(cè)試效果如下:
努力分享優(yōu)質(zhì)的計(jì)算機(jī)視覺(jué)相關(guān)內(nèi)容,歡迎關(guān)注:
個(gè)人微信(如果沒(méi)有備注不拉群!) 請(qǐng)注明:地區(qū)+學(xué)校/企業(yè)+研究方向+昵稱
下載1:何愷明頂會(huì)分享
在「AI算法與圖像處理」公眾號(hào)后臺(tái)回復(fù):何愷明,即可下載。總共有6份PDF,涉及 ResNet、Mask RCNN等經(jīng)典工作的總結(jié)分析
下載2:終身受益的編程指南:Google編程風(fēng)格指南
在「AI算法與圖像處理」公眾號(hào)后臺(tái)回復(fù):c++,即可下載。歷經(jīng)十年考驗(yàn),最權(quán)威的編程規(guī)范!
下載3 CVPR2021 在「AI算法與圖像處理」公眾號(hào)后臺(tái)回復(fù):CVPR,即可下載1467篇CVPR 2020論文 和 CVPR 2021 最新論文
點(diǎn)亮
,告訴大家你也在看
