熱影像物件偵測 (Object Detection on Thermal Image) - 1 使用TensorRT C++ Win10 - 測試model及匯出onnx model



最近疫情的關係熱影像攝影機的應用瞬間發展很快,所以想說有能不能以熱影像畫面為基礎,作物件辨識。在網路上尋找到這篇文章,使用yolov3-spp模型架構訓練,model細節及介紹請看原文,這裡主要是要說明我怎麼將他的model轉換成tensorRT可以執行的部分。想研究如果用c++執行這個model能不能更快一些。

一、clone github

1.clone github專案下來
2.Download pre-trained weights

我下載下來的 pre-trained weights 長這樣,根本不知道要使用哪個才對,如果有人嘗試其他的可以留言跟我說喔。

二、測試程式可以執行

專案根目錄裡有個detect.py,這是我們開始的主程式,我嘗試出以下組合是可以使用的。
類別名稱檔:data/coco.data
model_config:cfg/yolov3-spp.cfg
model:weights/yolov3-spp.weights

如果在data資料夾裡沒有找到 coco.data、coco.names檔案的話請解壓縮labels.zip

帶入所需的參數執行:

--data <類別名稱檔>
--cfg <model_config>
--weights <model>
--source <要辨識的圖片資料夾或影片檔>

python detect.py --data data/coco.data --cfg cfg/yolov3-spp.cfg --weights weights/weights/yolov3-spp.weights --source data/samples

這是他deteset裡面的一張圖的便是結果:


三、執行速度

我使用的系統環境是:
CPU:i7-10700
GPU:GTX1660 SUPER
CUDA:10.2
python:3.6.8
pytorch:1.8.1
onnx:1.8.0

每一張處理時間是:0.02秒左右

四、model匯出成onnx格式

1.切換匯出onnx model模式

在models.py 裡將 ONNX_EXPORT = False 更改為 True

2.設定model input image shape

在detect.py 裡改變 img_size ,我改成依照他opt.img_size預設值 416,但是需要高跟寬兩個參數,我將寬設定為416,高設定為320。
如果不修改img_size,他預設匯出model的input shape [1,3,320,192],格式[B,C,H,W]
B - batch size
C - number of channels
H - image height
W - image width

加入的地方如下圖:

3.加入onnx model input output name

在 Export mode 地方加入:
input_names = ["input0"]
output_names = ["output0", "output1"]

修改 torch.onnx.export() 內參數:
torch.onnx.export(model, img,'weights/export.onnx', verbose=False, opset_version=11)
改為
torch.onnx.export(model, img,'weights/export.onnx', verbose=False, opset_version=11,
                         input_names=input_names, output_names=output_names)

加入的地方如下圖:


4.匯出onnx model

python detect.py --data data/coco.data --cfg cfg/yolov3-spp.cfg --weights weights/weights/yolov3-spp.weights --source data/samples

在 weights 的資料夾裡面看到匯出的model: export.onnx

五、查看模型架構

python 安裝 netron 或是搜尋 netron 上傳model
model 的下半部大致上會長這樣:


六、參考資料來源:

https://medium.com/@joehoeller/object-detection-on-thermal-images-f9526237686a
https://github.com/joehoeller/Object-Detection-on-Thermal-Images

留言

這個網誌中的熱門文章

C# 模擬鍵盤滑鼠控制電腦

python nn 聲音辨識 -1 傅立葉轉換

android 定時通知(永久長期的) 本篇只講AlarmManager使用

python opencv 基本讀取、轉換、顯示、儲存等

python pyautogui 簡介