From 08126b175d65acd2f6ea9d45a0276246e1f954d0 Mon Sep 17 00:00:00 2001 From: wang-xinyu Date: Fri, 25 Sep 2020 17:05:32 +0800 Subject: [PATCH] fix yolov5 --- yolov5/common.hpp | 32 ++++++++--------- yolov5/yololayer.cu | 88 ++++++++++++++------------------------------- yolov5/yololayer.h | 1 + yolov5/yolov5.cpp | 9 ++--- 4 files changed, 44 insertions(+), 86 deletions(-) diff --git a/yolov5/common.hpp b/yolov5/common.hpp index c42c8cf..47806c1 100644 --- a/yolov5/common.hpp +++ b/yolov5/common.hpp @@ -32,8 +32,7 @@ cv::Mat preprocess_img(cv::Mat& img) { h = r_w * img.rows; x = 0; y = (Yolo::INPUT_H - h) / 2; - } - else { + } else { w = r_h * img.cols; h = Yolo::INPUT_H; x = (Yolo::INPUT_W - w) / 2; @@ -59,8 +58,7 @@ cv::Rect get_rect(cv::Mat& img, float bbox[4]) { r = r / r_w; t = t / r_w; b = b / r_w; - } - else { + } else { l = bbox[0] - bbox[2] / 2.f - (Yolo::INPUT_W - r_h * img.cols) / 2; r = bbox[0] + bbox[2] / 2.f - (Yolo::INPUT_W - r_h * img.cols) / 2; t = bbox[1] - bbox[3] / 2.f; @@ -299,7 +297,7 @@ int read_files_in_dir(const char *p_dir_name, std::vector &file_nam return 0; } -std::vector GetAnchors(std::map& weightMap) +std::vector getAnchors(std::map& weightMap) { std::vector anchors_yolo; Weights Yolo_Anchors = weightMap["model.24.anchor_grid"]; @@ -327,34 +325,32 @@ std::vector GetAnchors(std::map& weightMap) IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map& weightMap, IConvolutionLayer* det0, IConvolutionLayer* det1, IConvolutionLayer* det2) { auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1"); - std::vector anchors_yolo = GetAnchors(weightMap); + std::vector anchors_yolo = getAnchors(weightMap); PluginField pluginMultidata[4]; - int* NetData = new int[4]; + int NetData[4]; NetData[0] = Yolo::CLASS_NUM; NetData[1] = Yolo::INPUT_W; NetData[2] = Yolo::INPUT_H; NetData[3] = Yolo::MAX_OUTPUT_BBOX_COUNT; pluginMultidata[0].data = NetData; pluginMultidata[0].length = 3; - std::string name = "netdata"; - pluginMultidata[0].name = new char[name.size() + 1]; - strcpy(const_cast(pluginMultidata[0].name), name.c_str()); + pluginMultidata[0].name = "netdata"; pluginMultidata[0].type = PluginFieldType::kFLOAT32; int scale[3] = { 8, 16, 32 }; + int plugindata[3][8]; + std::string names[3]; for (int k = 1; k < 4; k++) { - int* plugindata = new int[8]; - plugindata[0] = Yolo::INPUT_W / scale[k - 1]; - plugindata[1] = Yolo::INPUT_H / scale[k - 1]; + plugindata[k - 1][0] = Yolo::INPUT_W / scale[k - 1]; + plugindata[k - 1][1] = Yolo::INPUT_H / scale[k - 1]; for (int i = 2; i < 8; i++) { - plugindata[i] = int(anchors_yolo[(k - 1) * 6 + i - 2]); + plugindata[k - 1][i] = int(anchors_yolo[(k - 1) * 6 + i - 2]); } - pluginMultidata[k].data = plugindata; + pluginMultidata[k].data = plugindata[k - 1]; pluginMultidata[k].length = 8; - std::string name = "yolodata" + std::to_string(k); - pluginMultidata[k].name = new char[name.size() + 1]; - strcpy(const_cast(pluginMultidata[k].name), name.c_str()); + names[k - 1] = "yolodata" + std::to_string(k); + pluginMultidata[k].name = names[k - 1].c_str(); pluginMultidata[k].type = PluginFieldType::kFLOAT32; } PluginFieldCollection pluginData; diff --git a/yolov5/yololayer.cu b/yolov5/yololayer.cu index f41e763..b3ed7c3 100644 --- a/yolov5/yololayer.cu +++ b/yolov5/yololayer.cu @@ -153,8 +153,6 @@ namespace nvinfer1 // Clone the plugin IPluginV2IOExt* YoloLayerPlugin::clone() const { - //YoloLayerPlugin *p = nullptr; - //p = new YoloLayerPlugin(); YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, mYoloKernel); p->setPluginNamespace(mPluginNamespace); return p; @@ -175,7 +173,6 @@ namespace nvinfer1 int info_len_i = 5 + classes; const float* curInput = input + bnIdx * (info_len_i * total_grid * CHECK_COUNT); - for (int k = 0; k < 3; ++k) { float box_prob = Logist(curInput[idx + k * info_len_i * total_grid + 4 * total_grid]); if (box_prob < IGNORE_THRESH) continue; @@ -198,18 +195,18 @@ namespace nvinfer1 int col = idx % yoloWidth; //Location - // pytorch： - // y = x[i].sigmoid() - // y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy - // y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh - //X: (sigmoid(tx) + cx)/FeaturemapW * netwidth + // pytorch: + // y = x[i].sigmoid() + // y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy + // y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh + // X: (sigmoid(tx) + cx)/FeaturemapW * netwidth det->bbox[0] = (col - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * netwidth / yoloWidth; det->bbox[1] = (row - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * netheight / yoloHeight; // W: (Pw * e^tw) / FeaturemapW * netwidth - // v5：https://github.com/ultralytics/yolov5/issues/471 + // v5: https://github.com/ultralytics/yolov5/issues/471 det->bbox[2] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 2 * total_grid]); - det->bbox[2] = det->bbox[2] * det->bbox[2] * anchors[2 * k]; // + det->bbox[2] = det->bbox[2] * det->bbox[2] * anchors[2 * k]; det->bbox[3] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 3 * total_grid]); det->bbox[3] = det->bbox[3] * det->bbox[3] * anchors[2 * k + 1]; det->conf = box_prob * max_cls_prob; @@ -272,66 +269,34 @@ namespace nvinfer1 IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) { - int ClassCount; - int YoloV5NetWidth; - int YoloV5NetHeight; - int MaxOutObject; - std::vector vYoloKernel; + int class_count = 80; + int input_w = 416; + int input_h = 416; + int max_output_object_count = 1000; + std::vector yolo_kernels(3); - int* floatdata0; - int* floatdata1; - int* floatdata2; - int* floatdata3; const PluginField* fields = fc->fields; for (int i = 0; i < fc->nbFields; i++) { - if (strcmp(fields[i].name, "netdata") == 0) { assert(fields[i].type == PluginFieldType::kFLOAT32); - floatdata0 = (int*)(fields[i].data); - ClassCount = floatdata0[0]; - YoloV5NetWidth = floatdata0[1]; - YoloV5NetHeight = floatdata0[2]; - MaxOutObject = floatdata0[3]; - } - if (strcmp(fields[i].name, "yolodata1") == 0) { + int *tmp = (int*)(fields[i].data); + class_count = tmp[0]; + input_w = tmp[1]; + input_h = tmp[2]; + max_output_object_count = tmp[3]; + } else if (strstr(fields[i].name, "yolodata") != NULL) { assert(fields[i].type == PluginFieldType::kFLOAT32); - floatdata1 = (int*)(fields[i].data); + int *tmp = (int*)(fields[i].data); YoloKernel kernel; - kernel.width = floatdata1[0]; - kernel.height = floatdata1[1]; - for (int j = 0; j < fields[i].length - 2; j++) - { - kernel.anchors[j] = floatdata1[j + 2]; + kernel.width = tmp[0]; + kernel.height = tmp[1]; + for (int j = 0; j < fields[i].length - 2; j++) { + kernel.anchors[j] = tmp[j + 2]; } - vYoloKernel.push_back(kernel); - } - if (strcmp(fields[i].name, "yolodata2") == 0) { - assert(fields[i].type == PluginFieldType::kFLOAT32); - floatdata2 = (int*)(fields[i].data); - YoloKernel kernel; - kernel.width = floatdata2[0]; - kernel.height = floatdata2[1]; - for (int j = 0; j < fields[i].length - 2; j++) - { - kernel.anchors[j] = floatdata2[j + 2]; - } - vYoloKernel.push_back(kernel); - } - if (strcmp(fields[i].name, "yolodata3") == 0) { - assert(fields[i].type == PluginFieldType::kFLOAT32); - floatdata3 = (int*)(fields[i].data); - YoloKernel kernel; - kernel.width = floatdata3[0]; - kernel.height = floatdata3[1]; - for (int j = 0; j < fields[i].length - 2; j++) - { - kernel.anchors[j] = floatdata3[j + 2]; - } - vYoloKernel.push_back(kernel); + yolo_kernels[2 - (fields[i].name[8] - '1')] = kernel; } } - std::reverse(vYoloKernel.begin(), vYoloKernel.end()); - YoloLayerPlugin* obj = new YoloLayerPlugin(ClassCount, YoloV5NetWidth, YoloV5NetHeight, MaxOutObject, vYoloKernel); + YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count, yolo_kernels); obj->setPluginNamespace(mNamespace.c_str()); return obj; } @@ -339,10 +304,9 @@ namespace nvinfer1 IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) { // This object will be deleted when the network is destroyed, which will - // call MishPlugin::destroy() + // call YoloLayerPlugin::destroy() YoloLayerPlugin* obj = new YoloLayerPlugin(serialData, serialLength); obj->setPluginNamespace(mNamespace.c_str()); return obj; } - REGISTER_TENSORRT_PLUGIN(YoloPluginCreator); } diff --git a/yolov5/yololayer.h b/yolov5/yololayer.h index a6d9a89..232e04e 100644 --- a/yolov5/yololayer.h +++ b/yolov5/yololayer.h @@ -131,6 +131,7 @@ namespace nvinfer1 static PluginFieldCollection mFC; static std::vector mPluginAttributes; }; + REGISTER_TENSORRT_PLUGIN(YoloPluginCreator); }; #endif diff --git a/yolov5/yolov5.cpp b/yolov5/yolov5.cpp index a39c17a..25d9b37 100644 --- a/yolov5/yolov5.cpp +++ b/yolov5/yolov5.cpp @@ -20,12 +20,11 @@ static const int INPUT_H = Yolo::INPUT_H; static const int INPUT_W = Yolo::INPUT_W; static const int CLASS_NUM = Yolo::CLASS_NUM; -static const int OUTPUT_SIZE = Yolo::MAX_OUTPUT_BBOX_COUNT * sizeof(Yolo::Detection) / sizeof(float) + 1; // we assume the yololayer outputs no more than 1000 boxes that conf >= 0.1 +static const int OUTPUT_SIZE = Yolo::MAX_OUTPUT_BBOX_COUNT * sizeof(Yolo::Detection) / sizeof(float) + 1; // we assume the yololayer outputs no more than MAX_OUTPUT_BBOX_COUNT boxes that conf >= 0.1 const char* INPUT_BLOB_NAME = "data"; const char* OUTPUT_BLOB_NAME = "prob"; static Logger gLogger; - // Creat the engine using only the API and not any parser. ICudaEngine* createEngine_s(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) { INetworkDefinition* network = builder->createNetworkV2(0U); @@ -431,8 +430,7 @@ int main(int argc, char** argv) { p.write(reinterpret_cast(modelStream->data()), modelStream->size()); modelStream->destroy(); return 0; - } - else if (argc == 3 && std::string(argv[1]) == "-d") { + } else if (argc == 3 && std::string(argv[1]) == "-d") { std::ifstream file(engine_name, std::ios::binary); if (file.good()) { file.seekg(0, file.end); @@ -443,8 +441,7 @@ int main(int argc, char** argv) { file.read(trtModelStream, size); file.close(); } - } - else { + } else { std::cerr << "arguments not right!" << std::endl; std::cerr << "./yolov5 -s // serialize model to plan file" << std::endl; std::cerr << "./yolov5 -d ../samples // deserialize plan file and run inference" << std::endl;