thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

GPU Batched NMS

Browse Source
This commit is contained in:
Marcos Luciano
2022-06-19 03:25:50 -03:00
parent 2300e3b44b
commit f621c0f429
24 changed files with 835 additions and 654 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp
View File

@@ -38,15 +38,15 @@ extern "C" bool NvDsInferParseYolo(
std::vector<NvDsInferParseObjectInfo>& objectList);
static NvDsInferParseObjectInfo convertBBox(
const float& bx, const float& by, const float& bw,
const float& bh, const uint& netW, const uint& netH)
const float& bx1, const float& by1, const float& bx2,
const float& by2, 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;
float x1 = bx1;
float y1 = by1;
float x2 = bx2;
float y2 = by2;
x1 = clamp(x1, 0, netW);
y1 = clamp(y1, 0, netH);
@@ -62,11 +62,11 @@ static NvDsInferParseObjectInfo convertBBox(
}
static void addBBoxProposal(
const float bx, const float by, const float bw, const float bh,
const float bx1, const float by1, const float bx2, const float by2,
const uint& netW, const uint& netH, const int maxIndex,
const float maxProb, std::vector<NvDsInferParseObjectInfo>& binfo)
{
NvDsInferParseObjectInfo bbi = convertBBox(bx, by, bw, bh, netW, netH);
NvDsInferParseObjectInfo bbi = convertBBox(bx1, by1, bx2, by2, netW, netH);
if (bbi.width < 1 || bbi.height < 1) return;
bbi.detectionConfidence = maxProb;
@@ -75,34 +75,25 @@ static void addBBoxProposal(
}
static std::vector<NvDsInferParseObjectInfo> decodeYoloTensor(
const float* detections,
const uint gridSizeW, const uint gridSizeH, const uint numBBoxes,
const uint numOutputClasses, const uint& netW, const uint& netH)
const int* counts, const float* boxes,
const float* scores, const float* classes,
const uint& netW, const uint& netH)
{
std::vector<NvDsInferParseObjectInfo> binfo;
for (uint y = 0; y < gridSizeH; ++y) {
for (uint x = 0; x < gridSizeW; ++x) {
for (uint b = 0; b < numBBoxes; ++b)
{
const int numGridCells = gridSizeH * gridSizeW;
const int bbindex = y * gridSizeW + x;
const float bx
= detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 0)];
const float by
= detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 1)];
const float bw
= detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 2)];
const float bh
= detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 3)];
const float maxProb
= detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 4)];
const int maxIndex
= detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 5)];
uint numBoxes = counts[0];
addBBoxProposal(bx, by, bw, bh, netW, netH, maxIndex, maxProb, binfo);
}
}
for (uint b = 0; b < numBoxes; ++b)
{
float bx1 = boxes[b * 4 + 0];
float by1 = boxes[b * 4 + 1];
float bx2 = boxes[b * 4 + 2];
float by2 = boxes[b * 4 + 3];
float maxProb = scores[b];
int maxIndex = classes[b];
addBBoxProposal(bx1, by1, bx2, by2, netW, netH, maxIndex, maxProb, binfo);
}
return binfo;
}
@@ -112,7 +103,6 @@ static bool NvDsInferParseCustomYolo(
NvDsInferNetworkInfo const& networkInfo,
NvDsInferParseDetectionParams const& detectionParams,
std::vector<NvDsInferParseObjectInfo>& objectList,
const uint &numBBoxes,
const uint &numClasses)
{
if (outputLayersInfo.empty())
@@ -130,18 +120,17 @@ static bool NvDsInferParseCustomYolo(
std::vector<NvDsInferParseObjectInfo> objects;
for (uint idx = 0; idx < outputLayersInfo.size(); ++idx)
for (uint idx = 0; idx < outputLayersInfo.size() / 4; ++idx)
{
const NvDsInferLayerInfo &layer = outputLayersInfo[idx];
assert(layer.inferDims.numDims == 3);
const uint gridSizeH = layer.inferDims.d[1];
const uint gridSizeW = layer.inferDims.d[2];
const NvDsInferLayerInfo &counts = outputLayersInfo[idx * 4 + 0];
const NvDsInferLayerInfo &boxes = outputLayersInfo[idx * 4 + 1];
const NvDsInferLayerInfo &scores = outputLayersInfo[idx * 4 + 2];
const NvDsInferLayerInfo &classes = outputLayersInfo[idx * 4 + 3];
std::vector<NvDsInferParseObjectInfo> outObjs =
decodeYoloTensor(
(const float*)(layer.buffer),
gridSizeW, gridSizeH, numBBoxes, numClasses,
(const int*)(counts.buffer), (const float*)(boxes.buffer),
(const float*)(scores.buffer), (const float*)(classes.buffer),
networkInfo.width, networkInfo.height);
objects.insert(objects.end(), outObjs.begin(), outObjs.end());
@@ -158,11 +147,10 @@ extern "C" bool NvDsInferParseYolo(
NvDsInferParseDetectionParams const& detectionParams,
std::vector<NvDsInferParseObjectInfo>& objectList)
{
uint numBBoxes = kNUM_BBOXES;
uint numClasses = kNUM_CLASSES;
int num_classes = kNUM_CLASSES;
return NvDsInferParseCustomYolo (
outputLayersInfo, networkInfo, detectionParams, objectList, numBBoxes, numClasses);
outputLayersInfo, networkInfo, detectionParams, objectList, num_classes);
}
CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseYolo);