乾貨 | YOLOV5 訓練自動駕駛數據集，並轉Tensorrt，收藏！

AI視線python

點擊關注上方「AI深度學習視線」，並「星標」公號
nginx

技術硬文，第一時間送達！git

準備數據集github
環境配置web
配置文件修改算法
訓練sql
推理shell
轉Tensorrtjson

1微信

準備數據集

1.1 BDD數據集

BDD100K是最大的開放式駕駛視頻數據集之一，其中包含10萬個視頻和10個任務，目的是方便評估自動駕駛圖像識別算法的的進展。每一個高分辨率視頻一共40秒。該數據集包括超過1000個小時的駕駛數據，總共超過1億幀。這些視頻帶有GPU / IMU數據以獲取軌跡信息。該數據集具備地理，環境和天氣多樣性，從而能讓模型可以識別多種場景，具有更多的泛化能力。這些豐富的戶外場景和複雜的車輛運動使感知任務更具挑戰性。該數據集上的任務包括圖像標記，車道檢測，可駕駛區域分割，道路對象檢測，語義分割，實例分割，多對象檢測跟蹤，多對象分割跟蹤，領域自適應和模仿學習。咱們能夠在BDD100K數據網站上下載數據。

Bdd100k的標籤是由Scalabel生成的JSON格式。
- labels [ ]: - id: int32 - category: string (classification) - manualShape: boolean (whether the shape of the label is created or modified manually) - manualAttributes: boolean (whether the attribute of the label is created or modified manually) - score: float (the confidence or some other ways of measuring the quality of the label.) - attributes: - occluded: boolean - truncated: boolean - trafficLightColor: "red|green|yellow|none" - areaType: "direct | alternative" (for driving area) - laneDirection: "parallel|vertical" (for lanes) - laneStyle: "solid | dashed" (for lanes) - laneTypes: (for lanes) - box2d: - x1: float - y1: float - x2: float - y2: float

道路對象類別包括如下幾類：

[ "bike", "bus", "car", "motor", "person", "rider", "traffic light", "traffic sign", "train", "truck"]

1.2 YOLO數據格式

每一個圖片文件.jpg，都有同一命名的標籤文件.txt。

標籤文件中每一個對象獨佔一行，格式爲<object-class> <x> <y> <width> <height>。

其中：

<object-class>-表示對象的類別序號：從0 到 (classes-1)
<x> <y> <width> <height> -參照圖片寬度和高度的相對比例(浮點數值)，從0.0到1.0
例如：<x> = <absolute_x> / <image_width>或<height> = <absolute_height> / <image_height>
注意：<x> <y>是矩形的中心，而不是左上角位置。

以下圖所示：

YOLO V5的標籤文件和圖像文件應位於同一目錄下。

1.3 BDD數據轉YOLO格式

Berkerley 提供了Bdd100k數據集的標籤查看及標籤格式轉化工具。因爲沒有直接從bdd100k轉換成YOLO的工具，所以咱們首先得使用將bdd100k的標籤轉換爲coco格式，而後再將coco格式轉換爲yolo格式。

bdd to coco

個人目的是識別包括不一樣顏色交通燈在內的全部交通對象，所以咱們須要對原版的bdd2coco.py進行一些修改，以獲取交通燈顏色併產生新的類別。

這是修改完的核心代碼：

for label in i['labels']: annotation = dict() category=label['category'] if (category == "traffic light"): color = label['attributes']['trafficLightColor'] category = "tl_" + color if category in id_dict.keys(): empty_image = False annotation["iscrowd"] = 0 annotation["image_id"] = image['id'] x1 = label['box2d']['x1'] y1 = label['box2d']['y1'] x2 = label['box2d']['x2'] y2 = label['box2d']['y2'] annotation['bbox'] = [x1, y1, x2-x1, y2-y1] annotation['area'] = float((x2 - x1) * (y2 - y1)) annotation['category_id'] = id_dict[category] annotation['ignore'] = 0 annotation['id'] = label['id'] annotation['segmentation'] = [[x1, y1, x1, y2, x2, y2, x2, y1]] annotations.append(annotation)

在完成bdd100k格式到yolo格式的轉換後，會得到bdd100k_labels_images_det_coco_train.json和bdd100k_labels_images_det_coco_val.json兩個文件。

Coco to yolo

在完成先前的轉換以後，咱們須要將訓練集和驗證集的coco格式標籤轉換爲yolo格式。注意須要分別指定訓練集和驗證集圖片位置，對應的coco標籤文件位置，及生成yolo標籤的目標位置。

config_train ={ "datasets": "COCO", "img_path": "bdd100k_images/bdd100k/images/100k/train", "label": "labels/bdd100k_labels_images_det_coco_train.json", "img_type": ".jpg", "manipast_path": "./", "output_path": "labels/trains/", "cls_list": "bdd100k.names", } config_valid ={ "datasets": "COCO", "img_path": "bdd100k_images/bdd100k/images/100k/val", "label": "labels/bdd100k_labels_images_det_coco_val.json", "img_type": ".jpg", "manipast_path": "./", "output_path": "labels/valids/", "cls_list": "bdd100k.names", }

除此以外，咱們還得將全部的類別寫入bdd100k.names文件。

personridercarbustruckbikemotortl_greentl_redtl_yellowtl_nonetraffic signtraintl_green

運行Bdd_preprocessing中的完整代碼能夠完成Bdd100k格式標籤到YOLO標籤格式的轉換。

Bdd2coco以及coco2yolo的詳細說明能夠參看bdd100k代碼庫和convert2Yolo代碼庫。

環境配置

2.1 官方代碼

https://github.com/ultralytics/yolov5/tree/v3.0

因爲後面轉tensorrt版本支持yolov5到3.0版本，因此以3.0版本進行實驗。

環境配置可經過下面命令進行一鍵配置。


# pip install -r requirements.txt
# base ----------------------------------------Cythonmatplotlib>=3.2.2numpy>=1.18.5opencv-python>=4.1.2pillowPyYAML>=5.3scipy>=1.4.1tensorboard>=2.2torch>=1.6.0torchvision>=0.7.0tqdm>=4.41.0
# coco ----------------------------------------# pycocotools>=2.0
# export --------------------------------------# packaging # for coremltools# coremltools==4.0b4# onnx>=1.7.0# scikit-learn==0.19.2 # for coreml quantization
# extras --------------------------------------# thop # FLOPS computation# seaborn # plotting

配置文件修改

3.1 修改 ./data/coco.yaml

修改./data/coco.yaml--》存爲bdd.yaml

修改內容：

（1）train/val/test 路徑

其中的txt內容均爲各集合圖像實際絕對路徑。

（2）nc：number class 類別數量，BDD數據類別爲10

（3）names：前面bdd數據集介紹時候已經列出

3.2 修改 ./model/yolov5.yaml

修改：nc爲BDD數據類別數：10

3.3 修改./train.py

修改：

（1）--weights，這裏s/m/l/x四個型號能夠選擇

（2）--cfg，這裏s/m/l/x四個型號能夠選擇

（3）--data，選擇根據coco.yaml修改後的bdd.yaml

（4）--batch-size 和 --img-size 能夠再這裏修改也能夠默認不動，再訓練命令行裏設定

訓練

預訓練模型

python train.py --img 640 --batch 32 --epochs 300 --data './data/bdd.yaml' --cfg ./models/custom_yolov5x.yaml --weights "./weights/yolov5x.pt" --name yolov5x_bdd_prew  --cache

從頭訓練

python train.py --img 640 --batch 32 --epochs 300 --data './data/bdd.yaml' --cfg ./models/custom_yolov5x.yaml --weights "" --name yolov5x_bdd  --cache

train_loss:

val_loss:

推斷

可選參數：

— weights: 訓練權重的路徑
— source：推理目標的路徑，能夠是圖片，視頻，網絡攝像頭等
— source：推理結果的輸出路徑
— img-size：推理圖片的大小
— conf-thres：對象置信閾值，默認0.4
— iou-thres：NMS的IOU閾值，能夠根據實際對象的重疊度調節，默認0.5
— device: 選擇使用CUDA或者CPU
— view-img：顯示全部推理結果
— save-txt：將每一幀的推理結果及邊界框的位置，存入*.txt文件
— classes：類別過濾，意思是隻推理目標類別
— agnostic-nms：使用agnostic-nms NMS

python detect.py --source 0 # webcam file.jpg # image file.mp4 # video path/ # directory path/*.jpg # glob rtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa # rtsp stream rtmp://192.168.1.105/live/test # rtmp stream                            http://112.50.243.8/PLTV/88888888/224/3221225900/1.m3u8  # http stream

轉Tensorrt

6.1 工程配置

https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5,

該項目提供了一大批常見模型的轉Tensorrt方法。

環境要求：

GTX1080 / Ubuntu16.04 / cuda10.0 / cudnn7.6.5 / tensorrt7.0.0 / nvinfer7.0.0 / opencv3.3

高版本tensorrt7的變化以下：

6.2 生成轉tensorrt的中間文件 yolov5.wts

拷貝 ./tensorrt/yolov5/gen_wts.py文件到./yolov5 工程下，修改其中加載模型路徑，執行該python文件，獲得yolov5.wts，並將其拷貝回 ./tensorrt/yolov5/下。

1. generate yolov5s.wts from pytorch with yolov5s.pt
git clone https://github.com/wang-xinyu/tensorrtx.gitgit clone https://github.com/ultralytics/yolov5.git// download its weights 'yolov5s.pt'// copy tensorrtx/yolov5/gen_wts.py into ultralytics/yolov5// ensure the file name is yolov5s.pt and yolov5s.wts in gen_wts.py// go to ultralytics/yolov5python gen_wts.py// a file 'yolov5s.wts' will be generated.

6.3 編譯yolov5並生成tensorrt模型yolov5.engine

編譯以前須要修改：

（1）選模型

（2）CMakeLists.txt

若是tensorrt是經過tar包解壓安裝的，還須要在CMakeList.txt中對tensorrt路徑進行指定，否則會報錯找不到nvinfer

（3）另外，若是系統是Ubuntu18.04的話還會存在opencv的問題，找不到libpng12.so和libjasper.so.

這個問題可經過https://blog.csdn.net/baobei0112/article/details/108991915 該博客內容找到答案。

（4）./tensorrt/yolov5/下新建個samples文件夾，把須要測試的圖片放進去。

作好準備工做，下面就能夠進行YOLOV5的engine編譯工做。

build tensorrtx/yolov5 and run
// put yolov5s.wts into tensorrtx/yolov5// go to tensorrtx/yolov5// ensure the macro NET in yolov5.cpp is smkdir buildcd buildcmake ..makesudo ./yolov5 -s // serialize model to plan file i.e. 'yolov5s.engine'sudo ./yolov5 -d  ../samples // deserialize plan file and run inference, the images in samples will be processed.

6.4 Tensorrt各yolo模型對比

潮水自會來去，但心志得堅若磐石。即使成不了那根定海神針，也至少不是那隨意被拍上岸的野鬼遊魂。

長按關注 "AI深度視線"
最新人工智能、深度學習、SLAM乾貨奉上！

本文分享自微信公衆號 - AI深度學習視線（AI_DeepSight）。
若有侵權，請聯繫 support@oschina.cn 刪除。
本文參與「OSC源創計劃」，歡迎正在閱讀的你也加入，一塊兒分享。