YOLOv5 實現無人機識別
https://mp.weixin.qq.com/s/700WcJJ1S--Qy6KAeSmmDw
點擊上方藍字關注我們
微信公眾號：OpenCV學堂
關注獲取更多計算機視覺與深度學習知識
編輯: gloomyfish
作者: 陳金光
引言
YOLOv5發布至今，已經迭代多個版本，是唯一基於pytorch框架的YOLO系列的模型框。簡單易用，特別適合自定義目標檢測與對象檢測。只要用了Pytorch版本YOLOv5框架，可以毫不誇張的說兩個小時就可以學會目標檢測。
01
第一步數據准備
想用視覺識別一下空中飛行的無人機，識別對象有：
1：鳥類bird2：無人機UAV3：直升機helicopter
想通過自己創建一個數據集，訓練yolov5，在調用detect.py來識別一下效果。
百度了一下，發現無法獲取別人已經標注好的數據集。（這里提醒大家，自己做數據集不費時間，給100多張圖打標簽label就行，用不了半個小時）圖片就從百度搜索就行，我們的目的是感性的認識這套牛掰的算法，100張圖訓練個幾個小時就可以出來很好的效果，還有不用非要顯卡，雖然我的電腦有顯卡，但是裝pytorch的顯卡cuda版本的總是出錯，索性不弄了，直接用cpu版本，cpu版本的pytorch包還小，一會兒就能安裝好，這100張圖我訓練了100個epoch花費了3到4個小時，我的電腦的cpu是i7-9700的）
我是從百度上找了大概130多張圖片，我截個圖：
用labelimage打上標簽，所有的標簽數量大概也就200來個吧！不會用labelimgage制作數據集的同學可以看這里：
https://github.com/tzutalin/labelImgyolo的數據文件夾格式又要鞏固一下了，先建一個數據文件夾data，在數據文件夾data目錄下面建圖像文件夾images和標簽文件夾labels，然后在images文件夾下面建立訓練數據集train和驗證集val（validation的縮寫，驗證校驗的意思），在labels的文件夾下面也是建立train和val兩個文件夾，分別放訓練集和驗證集的標簽。
另外說一下驗證集，驗證集里面我找了26張圖片，標記了40個左右的labels吧，所以所有的數據加起來大概150張圖片。
我放一個圖片看一下驗證集里面的圖片如下圖：
到這里為止，所有的數據基本整理完畢，另外還需要幾張測試的圖片，到測試的時候再從網上找就是了。
02
配置與修改參數
這里下載yolov5的源碼就不寫怎么下載了，自己百度就好，相對難度低，我是從gitee上下載的，github總是卡，這個源碼也沒有多大，幾分鍾的時間就下載好了。下面直接寫怎么改一下代碼訓練自己制作好的無人機的數據集。
分步驟寫吧
首先：明確思路
怎么訓練這個算法？我要訓練什么模型？我要拿這個模型做什么？這些問題是在訓練這個算法的時候首先要想的問題。
回答：
我要訓練一個模型能夠識別空中飛行的無人機，並打上紅框label，空中的復雜環境下可能不只有無人機，會有鳥、樹、白雲、直升機、飛機等等其他干擾因素印象，簡單起見，我又加上了鳥bird和直升機helicopter，所以用yolov5分類nc（number class）為3個，也就是無人機UAV、鳥bird、直升機、helicopter，看下圖，建立一個數據信息文件，文件格式yaml
我把這個文件名保存成了chenchen.yaml
訓練集圖片train在uav_train下面的data里面，訓練集也是，這個文件不用寫標簽label的地址，因為文件夾結構是定的，程序會自動判斷出label的位置。
chenchen.yaml文件里面的nc改成2，names改成自己的，也就是uav bird 和helicopter
下面我貼兩張圖片，這兩張圖片是我訓練了100張圖后，測試的標簽和測試的預測，基本上框的位置已經很好了，只有左上角的這個直升機helicopter沒有被識別出來，這個我感覺是因為我找的直升機圖片的問題，這個直升機圖片有點特別，訓練數據中這樣的圖片少的原因。
總結：思路明確了，就是做一個三分類模型，無人機、鳥、直升機，
其次修改train.py文件
思路確定了后就是訓練這個模型了，我們已經把數據集准備好了，現在就訓練這個模型。下載下來的yolo源碼里面的train.py文件，我截圖一下：
這個是train.py文件
第一：導入各種依賴包；第二：定義train函數方法，主要的就是這個train方法；第三：主程序；
解釋一下，各種依賴包不用解釋，train方法就是作者的核心了，簡單解釋就是yolo算法的原理，主程序里面就是調用tran方法畫框框，（到目前位置我能給的解釋也只有這么多，我沒有花費時間研究，我感性上的理解應該是這樣）
這個文件中主要要修改的參數是什么呢？
下面這張圖中我給出了箭頭
第一個箭頭：
yolov5s.pt，這個是權重文件，用的是s版本，small，yolov5的權重文件不同分為幾個不同版本，這里我用的是s版本，這里可以認為是模型的初始權重，
第二個箭頭：
這個yolov5s.yaml文件我保存到了uav_train文件夾里面，這個文件是配置yolov5模型的，叫做配置文件吧，等下這個文件我截圖一下，這個文件也要改一個地方。
第三個箭頭：
就是數據文件，就是上面定義的一個chenchen.yaml文件，這里有訓練數據集的地址。
第四個箭頭：
是epoch，我寫的是100，大家也可以改成10試一下，20分鍾訓練完成，我當時做了一下，10個epoch效果很差，100個的時候效果就是我視頻的效果，都能到70%准確率以上。
第五個箭頭：
是batch，我是16張圖片一個batch
第二個箭頭解釋圖片如下：
這里是這個算法的配置文件，最終輸出的類別是3，這改成3保存就行，其他的能看出來這個里面是yolo的架構，我也不會解釋，有想了解的可以自己百度。
總結：到這里train.py就修改完畢了
03
訓練與測試
到這里就可以訓練了，我用的是pycharm這個集成環境，所以選中train.py文件直接右鍵運行就可以
下面就是漫長的訓練時間，只有等着
其實這一步可以不用，test.py文件主要是用來測試驗證集val的，這里和訓練一樣，不要把訓練好的權重文件的地址換一下，就是best.pt文件，data的參數改一下，改成chenchen.yaml就可以。最激動人心的是第五步，就是訓練好的模型驗證一下其他的圖片，看看效果怎么樣。
測試新圖片，新視頻看效果
這個文件是detect.py，我截圖看下：
這里有四個箭頭
第一個：訓練好的最好的權重文件，加載上第二個：我從百度找了一個視頻文件保存成了666.mp4第三個：置信度，我這里其實可以寫0.7，第四個：iou，
解釋：如果訓練的epoch少，可以把三四箭頭改小點，要不然視頻和圖片中不會被標記出，會讓你誤認為自己訓練的不對。測試效果如