haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix yolov5

Browse Source
This commit is contained in:
wang-xinyu 2020-09-25 17:05:32 +08:00
parent 8bdbb47dfa
commit 08126b175d
4 changed files with 44 additions and 86 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
yolov5/common.hpp
View File

@ -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<std::string> &file_nam
return 0;
}
std::vector<float> GetAnchors(std::map<std::string, Weights>& weightMap)
std::vector<float> getAnchors(std::map<std::string, Weights>& weightMap)
{
std::vector<float> anchors_yolo;
Weights Yolo_Anchors = weightMap["model.24.anchor_grid"];
@ -327,34 +325,32 @@ std::vector<float> GetAnchors(std::map<std::string, Weights>& weightMap)
IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, IConvolutionLayer* det0, IConvolutionLayer* det1, IConvolutionLayer* det2)
{
auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1");
std::vector<float> anchors_yolo = GetAnchors(weightMap);
std::vector<float> 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<char*>(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<char*>(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;

88
yolov5/yololayer.cu
View File

@ -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<Yolo::YoloKernel> vYoloKernel;
int class_count = 80;
int input_w = 416;
int input_h = 416;
int max_output_object_count = 1000;
std::vector<Yolo::YoloKernel> 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);
}

1
yolov5/yololayer.h
View File

@ -131,6 +131,7 @@ namespace nvinfer1
static PluginFieldCollection mFC;
static std::vector<PluginField> mPluginAttributes;
};
REGISTER_TENSORRT_PLUGIN(YoloPluginCreator);
};
#endif

9
yolov5/yolov5.cpp
View File

@ -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<const char*>(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;