deepstream yolov5-6.0使用指南
首先修改 Yolov5-6.0的源码 yolo.py，将最后return，改为torch.cat(z, 1)
def forward(self, x):
z = [] # inference output
for i in range(self.nl):
x[i] = self.m[i](x[i]) # conv
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
if not self.training: # inference
if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
anchor_grid = (self.anchors[i].clone() * self.stride[i]).view(1, -1, 1, 1, 2)
y = x[i].sigmoid()
if self.inplace:
y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * anchor_grid # wh
else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (y[..., 2:4] * 2) ** 2 * anchor_grid # wh
y = torch.cat((xy, wh, y[..., 4:]), -1)
z.append(y.view(bs, -1, self.no))
return torch.cat(z, 1)
export导出的onnx是这样的：只有一个25200*6（1个类别）
然后在deepstream的/usr/src/tensorrt/bin下执行：onnx转engine
./trtexec --onnx=./personcartruck.onnx --saveEngine=./personcartruck.engine --fp16
后处理修改：
进入/opt/nvidia/deepstream/deepstream/sources/libs/nvdsinfer_customparser/nvdsinfer_custombboxparser.cpp
覆盖后执行，make install 会生成so到lib下，然后配置在config文件中，文章的最后有例子
/*
* Copyright (c) 2018-2020, NVIDIA CORPORATION. All rights reserved.
*
* Permission is hereby granted, free of charge, to any person obtaining a
* copy of this software and associated documentation files (the "Software"),
* to deal in the Software without restriction, including without limitation
* the rights to use, copy, modify, merge, publish, distribute, sublicense,
* and/or sell copies of the Software, and to permit persons to whom the
* Software is furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
* DEALINGS IN THE SOFTWARE.
*/
#include
#include
#include "nvdsinfer_custom_impl.h"
#include
#include
#include
#include
#define MIN(a,b) ((a) < (b) ? (a) : (b))
#define MAX(a,b) ((a) > (b) ? (a) : (b))
#define CLIP(a,min,max) (MAX(MIN(a, max), min))
#define DIVIDE_AND_ROUND_UP(a, b) ((a + b - 1) / b)
struct MrcnnRawDetection {
float y1, x1, y2, x2, class_id, score;
};
/* This is a sample bounding box parsing function for the sample Resnet10
* detector model provided with the SDK. */
/* C-linkage to prevent name-mangling */
extern "C"
bool NvDsInferParseYolo (std::vector const &outputLayersInfo,
NvDsInferNetworkInfo const &networkInfo,
NvDsInferParseDetectionParams const &detectionParams,
std::vector &objectList);
/* This is a sample bounding box parsing function for the tensorflow SSD models
* detector model provided with the SDK. */
static std::vector
nonMaximumSuppression(const float nmsThresh, std::vector binfo)
{
auto overlap1D = [](float x1min, float x1max, float x2min, float x2max) -> float {
if (x1min > x2min)
{
std::swap(x1min, x2min);
std::swap(x1max, x2max);
}
return x1max < x2min ? 0 : std::min(x1max, x2max) - x2min;
};
auto computeIoU
= [&overlap1D](NvDsInferParseObjectInfo& bbox1, NvDsInferParseObjectInfo& bbox2) -> float {
float overlapX
= overlap1D(bbox1.left, bbox1.left + bbox1.width, bbox2.left, bbox2.left + bbox2.width);
float overlapY
= overlap1D(bbox1.top, bbox1.top + bbox1.height, bbox2.top, bbox2.top + bbox2.height);
float area1 = (bbox1.width) * (bbox1.height);
float area2 = (bbox2.width) * (bbox2.height);
float overlap2D = overlapX * overlapY;
float u = area1 + area2 - overlap2D;
return u == 0 ? 0 : overlap2D / u;
};
std::stable_sort(binfo.begin(), binfo.end(),
[](const NvDsInferParseObjectInfo& b1, const NvDsInferParseObjectInfo& b2) {
return b1.detectionConfidence > b2.detectionConfidence;
});
std::vector out;
for (auto i : binfo)
{
bool keep = true;
for (auto j : out)
{
if (keep)
{
float overlap = computeIoU(i, j);
keep = overlap <= nmsThresh;
}
else
break;
}
if (keep) out.push_back(i);
}
return out;
}
float clamp(const float val, const float minVal, const float maxVal)
{
assert(minVal <= maxVal);
return std::min(maxVal, std::max(minVal, val));
}
static NvDsInferParseObjectInfo convertBBox(
const float& bx, const float& by, const float& bw,
const float& bh, const uint& netW, const uint& netH)
{
NvDsInferParseObjectInfo b;
float x1 = bx - bw / 2;
float y1 = by - bh / 2;
float x2 = x1 + bw;
float y2 = y1 + bh;
x1 = clamp(x1, 0, netW);
y1 = clamp(y1, 0, netH);
x2 = clamp(x2, 0, netW);
y2 = clamp(y2, 0, netH);
b.left = x1;
b.width = clamp(x2 - x1, 0, netW);
b.top = y1;
b.height = clamp(y2 - y1, 0, netH);
return b;
}
static void addBBoxProposal(
const float bx, const float by, const float bw, const float bh,
const uint& netW, const uint& netH, const int maxIndex,
const float maxProb, std::vector& binfo)
{
NvDsInferParseObjectInfo bbi = convertBBox(bx, by, bw, bh, netW, netH);
if (bbi.width < 1 || bbi.height < 1) return;
bbi.detectionConfidence = maxProb;
bbi.classId = maxIndex;
binfo.push_back(bbi);
}
static std::vector
nmsAllClasses(const float nmsThresh,
std::vector& binfo,
const uint numClasses)
{
std::vector result;
std::vector> splitBoxes(numClasses);
for (auto& box : binfo)
{
splitBoxes.at(box.classId).push_back(box);
}
for (auto& boxes : splitBoxes)
{
boxes = nonMaximumSuppression(nmsThresh, boxes);
result.insert(result.end(), boxes.begin(), boxes.end());
}
return result;
}
extern "C" bool NvDsInferParseYolo(
std::vector const& outputLayersInfo,
NvDsInferNetworkInfo const& networkInfo,
NvDsInferParseDetectionParams const& detectionParams,
std::vector& objectList)
{
const float kCONF_THRESH = detectionParams.perClassThreshold[0];
int dimensions = 5 + detectionParams.numClassesConfigured;
int confidenceIndex = 4;
int labelStartIndex = 5;
const float beta_nms = 0.45;
std::vector binfo;
for (unsigned int l = 0; l < 1; l++)
{
float* output = (float *)outputLayersInfo[l].buffer;
const NvDsInferLayerInfo &layer = outputLayersInfo[0];
//printf("zwh layer.inferDims.d[0] %d, layer.inferDims.d[1] %d \n", layer.inferDims.d[0],layer.inferDims.d[1]);
for (int i = 0; i < 25200; ++i) {
int index = i * dimensions; // index是5+classnum的倍数，位置是0；
if(output[index+confidenceIndex] <= 0.4f) continue;
for (int j = labelStartIndex; j < dimensions; ++j) {
output[index+j] = output[index+j] * output[index+confidenceIndex];
}
for (int k = labelStartIndex; k < dimensions; ++k) {
if(output[index+k] <= kCONF_THRESH) continue;
const float bx = output[index];
const float by = output[index+1];
const float bw = output[index+2];
const float bh = output[index+3];
const float maxProb = output[index+k];
const int maxIndex = k -5; // k是5+0的位置,所以类别是k-5
addBBoxProposal(bx, by, bw, bh, networkInfo.width, networkInfo.height, maxIndex, maxProb, binfo);
}
}
}
objectList.clear();
objectList = nmsAllClasses(beta_nms, binfo, detectionParams.numClassesConfigured);
return true;
}
/* Check that the custom function has been defined correctly */
CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseYolo);
然后在配置文件中配置：
#修改类别数
num-detected-classes=3
parse-bbox-func-name=NvDsInferParseYolo
custom-lib-path=/opt/nvidia/deepstream/deepstream-6.0/lib/libnvds_infercustomparser.so