基于YOLOv5的無人機目標圖像識別

張珺儀，郭 楠

（合肥工業大學 電氣與自動化工程學院，安徽 合肥 230009）

0 引 言

目前，直接使用YOLOv5官方權重進行無人機航拍目標檢測時，效果并不理想。無人機遙感場景下進行旋轉目標檢測數據規模小，大中小目標分布不均，背景復雜，目標尺度變化劇烈，小目標采樣易丟失，類間極度不平衡，導致其檢測概率小，精確度很低。同時，小目標（32×32）一直是目標檢測中的難題，例如檢測器在COCO上的小目標檢測平均準確率（Average Precision，AP）通常要低20個點[1-6]。本文將針對這類問題利用現有資源進行改進，選取新的DOTA數據集——DOTA-v2.0，并利用改進的YOLOv5算法對DOTA-v2.0數據集進行訓練，最終完成航拍目標的檢測。

1 DOTA數據集介紹

1.1 選擇DOTA數據集的理由

在一般對象數據集上，ImageNet和MSCOCO以圖像數量多、類別多、注釋詳細等特點受到研究人員的青睞。ImageNet在所有對象檢測數據集中擁有最多的圖像，然而，每張圖像的平均實例數遠小于MSCOCO和DOTA數據集。DOTA數據集中的圖像包含大量的對象實例，其中一些實例超過1 000個。與其他數據集比較，DOTA數據集擁有較大的平均示例數，如表1所示，相較于2018年發布的更龐大的DIOR數據集，更具備本文研究的包含小目標的航拍、遙感圖像特征，即在同一區域擁有較多的小目標對象。DOTA數據集與其他數據集的對比如表1所示[7]。

表1 DOTA數據集與其他數據集的對比

DOTA-v2.0作為最新的DOTA數據集，比DOTA-v1.0增加了航拍視角（Aerial）圖像。圖1是DOTA-v2.0和DOTA-v1.0的對比[8]。

圖1 DOTA數據集

由圖1可見，在圖均檢測邊框數上，DOTA-v2.0更加平均，且整體數目有所增加。綜上所述，本文選擇DOTA-v2.0為數據集。

在航拍圖像中，所使用的傳感器的分辨率和種類是產生數據集偏差的因素。為了消除偏差，數據集中的圖像是從多個傳感器和平臺（如谷歌地球）以多種分辨率收集的。為了增加數據的多樣性，本文收集了在多個城市拍攝的圖像。這些圖像由航空圖像解釋專家精心挑選。對于這些圖像，記錄每張圖片的位置和拍攝時間的準確地理坐標，以確保沒有重復的圖片。

DOTA-v2.0數據集共標注了18個類別，包括飛機、船舶、儲罐、棒球場、網球場、籃球場、地面田徑場、港口、橋梁、大型車輛、小型車輛、直升機、環形交叉路口、足球場、游泳池、集裝箱起重機、直升機停機坪以及機場[9]。

1.2 數據集結構

DOTA數據集的標注文件分為有向目標邊框（OBB）與水平目標邊框（HBB）兩個版本，其區別如圖2所示。HBB為無方向性的標注框，在小目標檢測時很容易造成邊框重疊，并且由于框內物體實際擁有朝向方向，容易增加檢測難度；而OBB選定了朝向方向，定向框更好地貼合了物體走向、形狀，減輕了計算量。根據本文關注的航拍、遙感小目標的密集性特點，本文選取OBB為標注方式。OBB標注的數據格式如表2所示。

表2 OBB標注格式

圖2 HBB標注與OBB標注的對比

2 訓練數據的改進與處理

處理數據時，本文在傳統增強方法的基礎上增加了馬賽克拼接法，即將多個待檢測圖像截取一部分合成為一張圖片進行整體檢測。這種方式能夠同時有效提升微小擾動和大量擾動條件下模型的檢測準確性。

2.1 網絡改進

遙感圖像由于分辨率過高，在訓練或預測過程中單張圖片運算時要求的顯存較高，成本要求高。對此，通常有以下兩種解決辦法。

（1）直接降低圖像的分辨率。將大于YOLOv5網絡標準圖片大小的圖像，直接調整圖像大小，例如將20 000×20 000大小的圖像直接調整至1 000×1 000，則大部分GPU都能滿足運算要求，但是很容易損失圖像信息，尤其對于DOTA-v1.5版本中新增加的許多小于10像素的示例極不友好，其尺寸會進一步縮小，進而可能致使其特征在深層網絡逐層提取中消失。

（2）切割單張圖像為多張圖像。即將大于YOLOv5網絡標準圖片大小的圖像，在送入程序前先進行分割處理，使其滿足要求。如將20 000×20 000的圖像裁剪成20×20=400張1 000×1 000大小的圖片，對切割的圖片進行多次預測，再將預測的結果進行拼接，同樣可以獲得原圖的預測結果。代價就是原本只需進行一次前向傳播，切割后需要進行100次前向傳播才能獲得預測結果，時間成本大幅增加。由于切割時常會導致處于切割后圖像邊緣部分的目標信息缺失，因此需要設置一定的重疊區域（gap）。由于DOTA-v2.0中的最大圖像大小超過20 000×20 000，為保證圖像中小目標的識別精度，取gap為切割圖像的50%。

本文選用第二種方法，并預先編寫分割圖像的函數，對圖像進行預處理。

2.2 label文本文件數據格式轉換

將分割后的label文本文件DOTA數據格式poly轉為YOLO數據格式xywh。

2.3 數據集的數據增強

為了提升模型的檢測效果，使其更好地適用于應用場景，同時增加訓練的數據量，提高模型的泛化能力，本文采用了圖像增強處理，即對數據集的數據增強。本節對數據集進行鏡像處理、隨機角度旋轉以及圖像平移等處理[10]。

經過垂直鏡像、水平鏡像、隨機角度旋轉處理后的結果示例如圖3所示。

圖3 數據增強處理結果

由圖3可見，經過垂直鏡像、水平鏡像、隨機角度旋轉處理后，圖像的方位、角度更具有隨機性，模型多樣性也得以提升。

2.4 檢測結果合并

利用被分割圖像檢測出的目標位置信息和圖像名稱中的裁剪位置信息，還原目標在原始未分割圖像中的位置。程序中考慮到由于裁剪圖像之間重疊區域的存在，圖像組合后會有一個目標被多個bbox標記，因此還在合并（merge）之后添加了一個poly_NMS。表3為分割前后的數據格式。

表3 分割前后的數據格式

通過訓練，模型在DOTA-V2.0數據集上的檢測結果如圖4所示，在DIOR數據集上的檢測結果如圖5所示。

由圖4和圖5可見，本文通過訓練得到的模型在DOTA-v2.0數據集和DIOR數據集上的檢測結果都很理想，說明了本模型具有普適性。

圖4 在DOTA-v2.0數據集上的檢測結果

圖5 在DIOR數據集上檢測結果

3 數據集訓練

由于數據集較大，因此本文中數據集采用線上GPU，型號為NVIDIA GeForce RTX 3090，epoch選為150，batchsize為4。數據集訓練過程截圖如圖6所示。

圖6 訓練過程截圖

4 數據集測試

數據集測試中，weights代表訓練的權重，source代表測試數據，可以是圖片/視頻路徑，也可以是‘0’（電腦自帶攝像頭），也可以是rtsp等視頻流；output代表網絡預測之后的圖片/視頻的保存路徑，img-size代表網絡輸入圖片大小，conf-thres代表置信度閾值，iou-thres代表做nms的iou閾值，device代表設置設備，view-img代表是否展示預測之后的圖片/視頻，默認False；save-txt代表是否將預測的框坐標以txt文件形式保存，默認False；classes代表設置只保留某一部分類別，例如設為class0或者class023；agnostic-nms代表進行nms是否將所有類別框一視同仁，默認為False；augment代表推理的時候進行多尺度、翻轉等操作（TTA）推理；update代表如果為True，則對所有模型進行strip_optimizer操作，去除pt文件中的優化器等信息，默認為False。

測試過程的部分程序如下所示。

經過程序運行，最終得到的測試結果有3 933張圖，截取其中1張，如圖7所示。

圖7 測試結果

5 模塊性能評估

對比原圖gt_poly和原圖detection_poly，求得各個類別的AP以及mAP，得出圖片包含類型、圖片名等，以評估模型性能。用detoutput表示detect.py函數的結果存放輸出路徑，imageset表示DOTA原圖數據集圖像路徑，annopath表示原始測試集的DOTAlabels路徑，classnames表示測試集中的目標種類部。評估結果如表2所示。

表2 檢測各類目標的平均準確率

由表2可見，用本文模型檢測目標的平均準確率能夠達到68%，部分類別的平均準確率可以達到90%。可見，該改進模型能滿足航拍目標的檢測精度要求，比原始訓練模型的精度有一定提高。

6 結 語

本文應用改進的YOLOv5算法對無人機航拍目標進行檢測，建立了基于YOLOv5算法的目標檢測模型，通過應用網絡改進、數據增強等技術，提高了模型檢測效能，完成了對航拍小目標的檢測，獲得了較高的檢測精度。