熱影像物件偵測 (Object Detection on Thermal Image) - 3 使用TensorRT C++ Win10 - 寫TensorRT C++
上篇文章我們測試了onnx格式的model，確認這個model是可以使用的。接下來就要在 TensorRT 上寫C++ 程式執行。
一、安裝tensorRT
可以參考我之前寫安裝tensorRT的文章，依照步驟安裝。
二、複製tensorRT sample 程式修改
我使用 samples\sampleOnnxMNIST 做修改，因為都是讀入onnx格式的model。
將 sampleOnnxMNIST 資料夾複製一份出來，一樣放在 samples
資料夾下，修改檔名叫做sampleOnnxThermalObjDetection，由於model input output
格式跟sample的不一樣，所以要稍微修改。
三、專案屬性加入opencv設定
我們需要讀入圖片，以及對圖片的處理。
這裡我借用openvino裡的opencv來用。
1.開啟專案的屬性頁面，在"其他 Include 目錄"裡添加 opencv include 路徑：
2.在"其他程式庫目錄"裡添加 opencv lib 路徑：
3.在"其他相依性"裡添加 opencv 使用到的 lib 檔名：
四、修改程式碼
開啟 sampleOnnxMNIST.cpp 檔案做修改
1. include opencv
#include <opencv2/core.hpp> #include <opencv2/highgui.hpp> #include <opencv2/imgcodecs.hpp> #include <opencv2/imgproc.hpp>
2.修改 model input output name：
將 initializeSampleParams() 修改成以下，之後要將model及image放在
samples/sampleOnnxThermalObjDetection/ 路徑之下。
samplesCommon::OnnxSampleParams initializeSampleParams(const samplesCommon::Args& args)
{
samplesCommon::OnnxSampleParams params;
if (args.dataDirs.empty()) //!< Use default directories if user hasn't provided directory paths
{
params.dataDirs.push_back("samples/sampleOnnxThermalObjDetection/");
}
else //!< Use the data directory provided by the user
{
params.dataDirs = args.dataDirs;
}
params.onnxFileName = "export.onnx";
params.inputTensorNames.push_back("input0");
params.outputTensorNames.push_back("output0");
params.outputTensorNames.push_back("output1");
params.dlaCore = args.useDLACore;
params.int8 = args.runInInt8;
params.fp16 = args.runInFp16;
return params;
}
3.修改input output get dimension的部分：
在宣告 mOutputDims 的地方再宣告一個 mOutputDims1。
在 build()
裡面的後半部在mEngine之後，看到 mInputDims及 mOutputDims的地方修改成以下：
assert(network->getNbInputs() == 1); mInputDims = network->getInput(0)->getDimensions(); assert(mInputDims.nbDims == 4); assert(network->getNbOutputs() == 2); mOutputDims = network->getOutput(0)->getDimensions(); assert(mOutputDims.nbDims == 2); mOutputDims1 = network->getOutput(1)->getDimensions(); assert(mOutputDims1.nbDims == 2);
4.修改 processInput()：
圖片 resize 成 model input大小，依照 [B,C,H,W] 的架構輸入，將每個 pixel
的顏色轉成浮點數，所以要除以 255.0。
bool SampleOnnxMNIST::processInput(const samplesCommon::BufferManager& buffers, cv::Mat &image)
{
const int inputChannels = mInputDims.d[1];
const int inputH = mInputDims.d[2];
const int inputW = mInputDims.d[3];
cv::Mat resized_image;
cv::resize(image, resized_image, cv::Size(inputW, inputH));
int batchIndex = 0;
int batchOffset = batchIndex * inputW * inputH * inputChannels;
float* hostDataBuffer = static_cast<float*>(buffers.getHostBuffer(mParams.inputTensorNames[0]));
// input shape [B,C,H,W]
for (size_t c = 0; c < inputChannels; c++) {
for (size_t h = 0; h < inputH; h++) {
for (size_t w = 0; w < inputW; w++) {
hostDataBuffer[batchOffset + c * inputW * inputH + h * inputW + w] =
float(float(resized_image.at<cv::Vec3b>(h, w)[c]) / 255.0);
}
}
}
return true;
}
5.修改 verifyOutput()：
- output0_class_score 輸出大小是[N,80]，80個種類每個種類的分數，所以要找出最高分的那個種類，N指的是有多少個物體。
- output1_bbox 輸出大小是[N,4]，物體位置四個標記點 [x,y,width,height]，而 x,y是物體的中心點位置，N指的是有多少個物體。
bool SampleOnnxMNIST::verifyOutput(const samplesCommon::BufferManager& buffers, cv::Mat image)
{
std::vector<cv::Rect2i> bbox_vec;
float* output0_class_score = static_cast<float*>(buffers.getHostBuffer(mParams.outputTensorNames[0])); // 80 class score
float* output1_bbox = static_cast<float*>(buffers.getHostBuffer(mParams.outputTensorNames[1])); // bbox
float conf_thres = 0.5;
const int output0_row = mOutputDims.d[0];
const int output0_col = mOutputDims.d[1];
const int output1_col = mOutputDims1.d[1];
for (int row = 0; row < output0_row; row++) {
float maxConfidence = output0_class_score[row * output0_col];
int classIdx = 0.0;
// get confidence highest score
for (int col = 1, first_idx = row * output0_col; col < output0_col; col++)
{
float confidence = output0_class_score[first_idx + col];
if (maxConfidence < confidence)
{
maxConfidence = confidence;
classIdx = col;
}
}
if (maxConfidence > conf_thres)
{
int idx = row * output1_col;
double x = double(output1_bbox[idx + 0]);
double y = double(output1_bbox[idx + 1]);
double w = double(output1_bbox[idx + 2]);
double h = double(output1_bbox[idx + 3]);
cv::Rect2i bbox;
bbox.x = static_cast<int>((x - w / 2) * float(image.cols));
bbox.y = static_cast<int>((y - h / 2) * float(image.rows));
bbox.width = static_cast<int>(w * float(image.cols));
bbox.height = static_cast<int>(h * float(image.rows));
if ((bbox.width > 15) && (bbox.height > 15))
{
bbox_vec.push_back(bbox);
}
}
}
// draw result
for (int idx = 0; idx < bbox_vec.size(); idx++)
{
cv::Rect2i i_rect = bbox_vec.at(idx);
cv::rectangle(image, i_rect, cv::Scalar(255, 0, 0), 2);
}
cv::imshow("img", image);
int key = cv::waitKey(0);
cv::destroyAllWindows();
return true;
}
6.infer()
讀入圖片，processInput()、verifyOutput() 輸入的參數增加輸入圖片進去。
bool SampleOnnxMNIST::infer()
{
// Create RAII buffer manager object
samplesCommon::BufferManager buffers(mEngine);
auto context = SampleUniquePtr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
if (!context)
{
return false;
}
// read image
cv::Mat image = cv::imread("image.jpeg");
if (image.cols == 0 || image.rows == 0)
{
printf("image is empty\n");
return false;
}
// Read the input data into the managed buffers
assert(mParams.inputTensorNames.size() == 1);
if (!processInput(buffers, image))
{
return false;
}
// Memcpy from host input buffers to device input buffers
buffers.copyInputToDevice();
bool status = context->executeV2(buffers.getDeviceBindings().data());
if (!status)
{
return false;
}
// Memcpy from device output buffers to host output buffers
buffers.copyOutputToHost();
// Verify results
if (!verifyOutput(buffers, image))
{
return false;
}
return true;
}
五、放入model及image檔案
將model及image放在 samples/sampleOnnxThermalObjDetection/ 路徑之下。
六、建置專案及執行結果
下圖為執行結果：
七、可改進的地方
會發現結果圖上有多個框重疊在同一個物體上，可以計算框與框之間的IOU值，來合併在同一個物體上的多個框。
八、參考資料來源：
https://medium.com/@joehoeller/object-detection-on-thermal-images-f9526237686a
https://github.com/joehoeller/Object-Detection-on-Thermal-Images
留言
張貼留言