基於垃圾目標檢測任務的YOLOv5初探

基於垃圾目標檢測任務的YOLOv5初探

做者：餘敏君

研究背景

垃圾分類做爲一種有效處理垃圾的科學管理方案，在提升資源利用率、緩解垃圾生產壓力以及改善生態環境等方面具備重要意義，是我國社會主義現代化和城市化進程中所必須採起的策略，備受世界各國的迫切關注。2019年以來，隨着上海市、杭州市等垃圾分類重點城市有關生活垃圾分類的立法、執法和監督等工做的順利開展，人們對垃圾分類相關話題的關注度日漸提高，我的垃圾分類的意識也有了很大的提升。但與此同時，因爲垃圾的種類極其豐富，我的對垃圾歸類的模糊程度廣泛較高，所以，垃圾分類自動化的實現顯得尤其重要。垃圾目標檢測做爲垃圾分類自動化的一個重要環節，本文將嘗試實現該過程。所謂目標檢測，簡單來說就是檢測圖像中的對象是什麼以及在哪裏的問題，即"目標分類+定位"問題。

YOLO系列模型做爲one-stage類目標檢測任務的表明算法，憑藉其快速、泛化性能好等特性深受研究者喜好。在前不久，YOLOv5也在GitHub上正式發佈，其當即在網上引起了普遍熱議。本文將參見官方提供的使用教程，嘗試簡單利用YOLOv5網絡模型在TACO數據集上實現垃圾目標檢測任務。

數據集處理

TACO是一個數據正在不斷增加的垃圾對象數據集，其以樹林、道路和海灘爲拍攝背景，目前包含60個類別的垃圾對象，共1500張圖像，近5千份標註。該數據集項目參見：https://github.com/pedropro/TACO

一方面，考慮到該數據集文件存放格式和標籤格式要符合YOLOv5模型的相關要求；另外一方面，也考慮到該數據集中各垃圾類型對象的樣本數量極其不均(如圖1所示)，所以，本文首先須要對數據集進行必要的操做。對於數據集的處理代碼可參見Mo項目[1]，項目中的_readme.ipynb文件詳細給出了相關代碼使用的具體操做。

圖1 TACO各種別對象的樣本數量

對於TACO數據集的處理主要可分爲如下兩個過程：

(1)因爲TACO提供的標籤爲coco類型(全部信息存放在annotations.json文件中)，項目首先須要將相關標籤轉換爲yolo類型。與此同時，考慮到本文僅對相關網絡模型進行初探和我的硬件設備的欠缺，該項目僅挑選出知足 要求的垃圾對象進行相關實驗。上述操做的相關代碼能夠參考博客[1]提供的代碼進行修改，該代碼能夠做爲模板用於知足自定義化coco格式轉yolo格式需求。

經初步統計，符合要求的垃圾類型(共8類)及其原編號以下所示：

'Clear plastic bottle': 5

'Plastic bottle cap': 7

'Drink can': 12

'Other plastic': 29

'Plastic film': 36

'Other plastic wrapper': 39

'Unlabeled litter': 58

'Cigarette': 59

對於coco格式轉yolo格式需求，其主要需完成如下三項任務：

· 存儲生成的標籤和圖像分別至兩個文件目錄中，同時標籤和圖像的文件命名要求一致

· 將目標對象原標籤集合遞增順序映射至{0-7}空間

· 因爲原始標籤中位置信息爲{top_x, top_y, width, height}，項目須要將其轉換爲{center_x, center_y, width, height}格式，並將其數值進行歸一化操做

其核心代碼部分以下所示(cocotoyolo.py)：

# 將垃圾類型原編號映射至{0-7}空間

label_transfer = {5: 0, 7: 1, 12: 2, 29: 3,

36: 4, 39: 5, 58: 6, 59: 7}

class_num = {} # 記錄各種型樣本數量

img_ids = data_source.getImgIds()

# 遍歷每張圖片，對標籤進行轉換

for index, img_id in tqdm.tqdm(enumerate(img_ids), desc='change .json file to .txt file'):

img_info = data_source.loadImgs(img_id)[0]

# 將含文件夾的路徑修改成文件名

save_name = img_info['file_name'].replace('/', '_')

# 移去文件擴展名

file_name = save_name.split('.')[0]

# 獲取單張圖像的寬和高

height = img_info['height']

width = img_info['width']

# 轉換所得txt文件存儲路徑

save_path = save_base_path + file_name + '.txt'

is_exist = False # 記錄圖片是否包含目標垃圾類型對象

with open(save_path, mode='w') as fp:

# 根據圖片編號找出垃圾對象的編號集合

annotation_id = data_source.getAnnIds(img_id)

boxes = np.zeros((0, 5))

if len(annotation_id) == 0: # 集合大小爲0

fp.write('')

continue

# 獲取coco格式的標籤

annotations = data_source.loadAnns(annotation_id)

lines = '' # 記錄轉換後yolo格式的標籤

# 遍歷對象標籤集

for annotation in annotations:

# 獲取垃圾對象的標籤

label = coco_labels_inverse[annotation['category_id']]

if label in label_transfer.keys():

# 垃圾類型屬於目標垃圾類型則進行格式轉換

is_exist = True

box = annotation['bbox']

if box[2] < 1 or box[3] < 1:

# 若是原標籤中出現無長或寬數據的狀況則跳過

continue

# top_x,top_y,width,height==>cen_x,cen_y,width,height

box[0] = round((box[0] + box[2] / 2) / width, 6)

box[1] = round((box[1] + box[3] / 2) / height, 6)

box[2] = round(box[2] / width, 6)

box[3] = round(box[3] / height, 6)

label = label_transfer[label] # 標籤映射

if label not in class_num.keys():

class_num[label] = 0

class_num[label] += 1

lines = lines + str(label) # 先存儲標籤

for i in box: # 再存儲位置信息

lines += ' ' + str(i)

lines += '\n' # 換行

fp.writelines(lines)

if is_exist:

# 存在目標類型對象，則拷貝圖像至指定目錄

shutil.copy('data/{}'.format(img_info['file_name']), os.path.join(save_image_path, save_name))

else:

# 不存在則刪除所生成的標籤文件

os.remove(save_path)

(2)完成標籤集的生成後，項目須要對其進行樣本劃分。首先，項目須要將樣本集按訓練集使用所有樣本，測試集樣本佔總樣本0.1比例的要求(其樣本圖片的數量分別爲1086和109，訓練集使用所有樣本是考慮到我的硬件設備較差的緣由)進行劃分。其次，考慮網絡模型對樣本存儲目錄的要求，項目須要將相應生成的圖像和標籤存儲至相應的文件夾下，其文件目錄格式如圖2所示。其實現代碼參見Mo項目中的sample.py文件。

圖2 數據集存儲文件目錄

模型配置與訓練

YOLOv5項目的下載地址：https://github.com/ultralytics

Mo平臺上部署的項目參見同名項目[2]，如下的全部內容能夠參考官方教程[2]。該Mo項目的具體操做流程請參見根目錄下的_readme.ipynb文件

模型配置

該部分主要涉及相關依賴安裝和配置文件設置這兩方面內容，接下來，本文將對上述內容進行簡要說明。

(1)依賴安裝

YOLOv5是由PyTorch深度學習框架搭建而成，所以，咱們首先須要在Python中安裝PyTorch框架，安裝教程能夠參見官網的相關內容。此處給出PyTorch最新CPU版本的安裝命令。

pip install torch==1.5.1+cpu torchvision==0.6.1+cpu -f https://download.pytorch.org/..._stable.html

除此以外，YOLOv5模型運行還須要安裝額外的第三方依賴包。官方已經將其所有放置在requirements.txt文件中。爲了安裝依賴儘量不出錯，我的對該文件提出如下兩點修改意見：

1. 本人fork該官方項目時，numpy的版本要求是1.17，但實際安裝時會出現一些問題。所以，我的建議將其版本要求修改成1.17.3。
2. 該依賴文件所給出的cocoapi.git下載的地址，在實際運行時下載的速度是很是緩慢的，並且在Windows系統下還可能出現異常。所以建議將相應語句用下列語句進行代替：

git+https://github.com/philferrie...

（2）配置文件設置

YOLOv5項目主要須要配置以下兩個配置文件：

1. data路徑下的taco.yaml配置文件（本身建立）。該文件主要配置訓練集和測試集存儲圖片的路徑以及對象類型信息。對於標籤存儲路徑的配置，項目會自動根據圖片存儲路徑來肯定，這也是爲何咱們的數據集文件目錄須要符合圖2所示要求的緣由。該配置文件的內容以下所示。

# train and val datasets (image directory or *.txt file with image paths)

train: taco/images/train/

val: taco/images/test/

# number of classes

nc: 8

# class names

names: ['Clear plastic bottle', 'Plastic bottle cap',

'Drink can',

'Other plastic',

'Plastic film', 'Other plastic wrapper',

'Unlabeled litter', 'Cigarette']

1. models路徑下與模型配置相關的配置文件編寫。根據官方教程[2]，研究者能夠直接修改models目錄下已存在的 yolov5_.yaml系列文件來進行相關配置。這裏，本人修改的是yolov5s.yaml文件(由於該實驗使用yolov5s.pt做爲模型預訓練權重)。我的只要將該文件中nc屬性的值修改成與上述taco.yaml的nc屬性值一致便可。

模型訓練

一方面，因爲我的電腦內存相對較小，直接使用官方提供的訓練參數設置將會致使內存爆炸而沒法工做。所以，該實驗不得不經過減小輸入圖像規模和batch_size值來促使模型進行訓練。另外一方面，因爲我的電腦GPU性能較差，該實驗選擇直接使用CPU來訓練相關模型，所以模型訓練速度相對較慢。考慮到上述兩方面的限制，本實驗模型訓練時所涉及的相關參數配置如表1所示：

表1 模型訓練相關的參數配置

命令行參數

參數含義

設置值

--img

統一輸入圖像規模

320

--batch

每次網絡訓練輸入圖像的數量

4

--epochs

整個數據集參與訓練的次數

100

--data

數據集配置文件路徑

./data/taco.yaml

--cfg

模型配置文件路徑

./models/yolov5s.yaml

--device

訓練的設備(CPU or GPU)

cpu

--weights

預訓練模型的權重文件

yolov5s.pt

模型訓練的調用命令以下所示：

python train.py --img 320 --batch 4 --epochs 100 --data ./data/taco.yaml --cfg ./models/yolov5s.yaml --device cpu --weights yolov5s.pt

關於預訓練權重的下載，這裏就不進行詳細介紹，百度一下應該能夠找到許多國內下載的資源。在這裏，本人在項目根目錄下放置了yolov5s的權重文件(yolov5s.pt)以方便研究者訓練模型。固然，咱們徹底能夠不使用預訓練權重來直接進行模型訓練，只要移去上述命令的--weights yolov5s.pt部分便可。

效果展現

根據官方提供的教程，模型訓練所生成的各種結果將被自動放置在根目錄的runs文件夾下。其中，weights文件夾下將會存放模型訓練所生成的效果最好和時間最近的權重文件，咱們能夠用這些文件完成模型的調用任務；results.txt文件存放着模型訓練過程當中的各項指標輸出，YOLO項目還自動對該輸出結果進行了可視化操做，生成了對應的圖表圖像。本實驗各項指標的輸出可視化圖像如圖3所示，上方五張爲訓練集對應結果，下方五張爲測試集對應結果。

圖3 模型訓練過程各項輸出指標的可視化圖像

利用detect.py文件咱們能夠利用生成的模型來檢測目標圖像中是否存在須要檢測的對象。爲了方便起見，咱們將生成的模型權重best.pt放置在項目根目錄下，將須要檢測的圖像inference/images文件夾下。參考官方文檔，咱們只需運行以下代碼便可展現模型目標檢測的效果：

python detect.py --weights best.pt --img 320 --conf 0.4

目標圖像檢測的效果如圖4所示，檢測生成的圖像位於inference/output文件夾中。其中，batch_1_000048圖像中錯誤檢測出了Drink can對象(此處主要與conf設置值的大小有關，我的能夠對其進行適當的調整)。

圖4 batch_1_000029(左)和batch_1_000048(右)的檢測效果

實驗小結

本文主要依託垃圾對象目標檢測任務，利用新穎的YOLOv5模型對TACO垃圾對象數據集進行了實驗。因爲我的硬件設備的不足和時間的限制，參考官方給出的一些輸出指標效果圖和自行實驗所得的圖像檢測結果，能夠看到本文最終所獲得的模型的性能其實相對來講不算好(YOLO系列模型自己就是須要跑很是久的時間)。感興趣地朋友能夠嘗試增大統一的圖像規模、增大batch_size和epoch等方式來提升該模型的目標檢測性能。

參考文獻

項目地址

[1]TACO：https://momodel.cn/workspace/5f0e734c95faedbb53ab3b26?type=app

[2]yolov5：https://momodel.cn/workspace/5f0e7c929fda75fe7f4b01e9?type=app

主要文獻

[1]將COCO中的.json文件轉成.txt文件：https://www.jianshu.com/p/8ddd8f3fdf73

[2]Train Custom Data：https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data