Nvidia Nano + Tensor RT + Tiny YOLOv3

資料來源: https://github.com/zombie0117/yolov3-tiny-onnx-TensorRT

不過有些參數要修一下 XD

這個code使用Python2, 前面的 TRT + YOLOv3 是用Python3, 剛好沒影響

git clone https://github.com/zombie0117/yolov3-tiny-onnx-TensorRT
cd yolov3-tiny-onnx-TensorRT

sudo apt-get install python-pip  #如果沒有pip再執行這個

# 原本的req.txt是寫onnx>=1.1.1, 可是後面又說onnx==1.4.1, 建議先修改這個txt檔
pip install -r requirements.txt

# 如果已經裝好onnx, 可以用下面指令刪掉重裝
pip uninstall onnx
pip install onnx==1.4.1

相關套件裝完後, 要先去改一下py裡面的參數. 主要是路徑的部分, 原本是寫絕對路徑, 可以照自己的環境修改 or 改成相對路徑

1. [yolov3-tiny.cfg]
檢查最後一行是不是空白的, 如果沒有的話要自己手動加一行空白進去
The cfg-file's last line must be a blank line. 
You should press Enter to add a blank line if there is no blank line at the end of the file.

2. [data_processing.py] 改路徑名稱
line14: LABEL_FILE_PATH = '/home/nvidia/yolov3-tiny2onnx2trt/coco_labels.txt'
line19: CATEGORY_NUM = 80

3. [yolov3_to_onnx.py] 改路徑名稱
line778: img_size = 416
line784: cfg_file_path = '/home/nvidia/yolov3-tiny2onnx2trt/yolov3-tiny.cfg'
line811: weights_file_path = '/home/nvidia/yolov3-tiny2onnx2trt/yolov3-tiny.weights'
line826: output_file_path = 'yolov3-tiny.onnx'

4. [onnx_to_tensorrt.py] 原本是imagelist.txt, 實際檔名是imageslist.txt, 記得一致 XD
還有改路徑名稱, output_shapes的語法也要改
line39: input_size = 416
line40: batch_size = 1
line42~line46:
    onnx_file_path = 'yolov3-tiny.onnx'
    engine_file_path = 'yolov3-tiny.trt'
    input_file_list = '/home/nvidia/yolov3-tiny2onnx2trt/imagelist.txt'
    IMAGE_PATH = '/home/nvidia/yolov3-tiny2onnx2trt/images/'
    save_path = '/home/nvidia/yolov3-tiny2onnx2trt/'

line 168:
    output_shapes = output_shapes_dic[str(input_size)]  # 這是原本的, 執行後會出現錯誤訊息. 可以改成下面這個
    output_shapes = [(batch_size, 255, 13, 13), (batch_size, 255, 26, 26)]  # 改成這樣可以跑出結果

PS. Line 168的來源: https://github.com/zombie0117/yolov3-tiny-onnx-TensorRT/issues/2

這些步驟準備完後, 應該就可以跑了

python yolov3_to_onnx.py    # 這個跑完會產生 yolov3-tiny.onnx 檔

python onnx_to_tensorrt.py  # 這個跑完後, 終端機會有辨識結果, 辨識完的圖片會儲存起來

那個常見的 dog.jpg 辨識結果:

Reading engine from file yolov3-tiny.trt
(416, 416)
time spent: 0.147424221039
2
test
boxes size: 3
[[123.64821648 219.95425227 258.68489429 295.94818745]
 [465.01741724  81.80695242 221.00351968  89.57543333]
 [530.07878114  93.9602292   99.6300466   67.47742823]] [0.81647941 0.68578115 0.628142  ] [16  2  7]
Saved image with bounding boxes of detected objects to save/dog_416_bboxes.png.