乾貨 | YOLOV5 訓練自動駕駛資料集，並轉成tensorrt【左側有碼】

關注公號‘AI深度學習視線’發表於2020-10-10
YOLO自動駕駛

 

 

  • 準備資料集
  • 環境配置
  • 配置檔案修改
  • 訓練
  • 推理
  • 轉Tensorrt
  • 遇到的Bugs

 

一、資料集準備

1，BDD資料集

讓我們來看看BDD100K資料集的概覽。

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"
]

我們實際關注的只有- labels [ ]欄目下的內容。

 

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的標籤檔案和影像檔案應位於同一目錄下。

 

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檔案。

person
rider
car
bus
truck
bike
motor
tl_green
tl_red
tl_yellow
tl_none
traffic sign
train
tl_green

執行Bdd_preprocessing中的完整程式碼可以完成Bdd100k格式標籤到YOLO標籤格式的轉換。

Bdd2coco以及coco2yolo的詳細說明可以參看bdd100k程式碼庫和convert2Yolo程式碼庫。

 

二、環境配置

1，官方程式碼

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

由於後面轉tensorrt版本支援yolov5到3.0版本，所以以3.0版本進行實驗。

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


# pip install -r requirements.txt

# base ----------------------------------------
Cython
matplotlib>=3.2.2
numpy>=1.18.5
opencv-python>=4.1.2
pillow
PyYAML>=5.3
scipy>=1.4.1
tensorboard>=2.2
torch>=1.6.0
torchvision>=0.7.0
tqdm>=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

三、配置檔案修改

1，修改./data/coco.yaml--》存為bdd.yaml

修改內容：

（1）train/val/test 路徑

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

（2）nc：number class 類別數量，BDD資料類別為10

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

 

2，./model/yolov5.yaml ：

修改：nc為BDD資料類別數：10

 

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 --data coco.yaml --cfg yolov5s.yaml --weights '' --batch-size 64
                                         yolov5m                                40
                                         yolov5l                                24
                                         yolov5x                                16

訓練過程中停止後 二次訓練：

有預訓練模型
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

 

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的變化如下：

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.git
git 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/yolov5
python gen_wts.py
// a file 'yolov5s.wts' will be generated.

 

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 s
mkdir build
cd build
cmake ..
make
sudo ./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.

4， Tensorrt各yolo模型對比

 

 

 

 

 

 

 

相關文章