Spaces:
Sleeping
Sleeping
File size: 3,713 Bytes
186701e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
#include "nvdsinfer_custom_impl.h"
#include <cassert>
#include <iostream>
/**
* Function expected by DeepStream for decoding the MMYOLO output.
*
* C-linkage [extern "C"] was written to prevent name-mangling. This function must return true after
* adding all bounding boxes to the objectList vector.
*
* @param [outputLayersInfo] std::vector of NvDsInferLayerInfo objects with information about the output layer.
* @param [networkInfo] NvDsInferNetworkInfo object with information about the MMYOLO network.
* @param [detectionParams] NvDsInferParseDetectionParams with information about some config params.
* @param [objectList] std::vector of NvDsInferParseObjectInfo objects to which bounding box information must
* be stored.
*
* @return true
*/
// This is just the function prototype. The definition is written at the end of the file.
extern "C" bool NvDsInferParseCustomMMYOLO(
std::vector<NvDsInferLayerInfo> const& outputLayersInfo,
NvDsInferNetworkInfo const& networkInfo,
NvDsInferParseDetectionParams const& detectionParams,
std::vector<NvDsInferParseObjectInfo>& objectList);
static __inline__ float clamp(float& val, float min, float max)
{
return val > min ? (val < max ? val : max) : min;
}
static std::vector<NvDsInferParseObjectInfo> decodeMMYoloTensor(
const int* num_dets,
const float* bboxes,
const float* scores,
const int* labels,
const float& conf_thres,
const unsigned int& img_w,
const unsigned int& img_h
)
{
std::vector<NvDsInferParseObjectInfo> bboxInfo;
size_t nums = num_dets[0];
for (size_t i = 0; i < nums; i++)
{
float score = scores[i];
if (score < conf_thres)continue;
float x0 = (bboxes[i * 4]);
float y0 = (bboxes[i * 4 + 1]);
float x1 = (bboxes[i * 4 + 2]);
float y1 = (bboxes[i * 4 + 3]);
x0 = clamp(x0, 0.f, img_w);
y0 = clamp(y0, 0.f, img_h);
x1 = clamp(x1, 0.f, img_w);
y1 = clamp(y1, 0.f, img_h);
NvDsInferParseObjectInfo obj;
obj.left = x0;
obj.top = y0;
obj.width = x1 - x0;
obj.height = y1 - y0;
obj.detectionConfidence = score;
obj.classId = labels[i];
bboxInfo.push_back(obj);
}
return bboxInfo;
}
/* C-linkage to prevent name-mangling */
extern "C" bool NvDsInferParseCustomMMYOLO(
std::vector<NvDsInferLayerInfo> const& outputLayersInfo,
NvDsInferNetworkInfo const& networkInfo,
NvDsInferParseDetectionParams const& detectionParams,
std::vector<NvDsInferParseObjectInfo>& objectList)
{
// Some assertions and error checking.
if (outputLayersInfo.empty() || outputLayersInfo.size() != 4)
{
std::cerr << "Could not find output layer in bbox parsing" << std::endl;
return false;
}
// Score threshold of bboxes.
const float conf_thres = detectionParams.perClassThreshold[0];
// Obtaining the output layer.
const NvDsInferLayerInfo& num_dets = outputLayersInfo[0];
const NvDsInferLayerInfo& bboxes = outputLayersInfo[1];
const NvDsInferLayerInfo& scores = outputLayersInfo[2];
const NvDsInferLayerInfo& labels = outputLayersInfo[3];
// num_dets(int) bboxes(float) scores(float) labels(int)
assert (num_dets.dims.numDims == 2);
assert (bboxes.dims.numDims == 3);
assert (scores.dims.numDims == 2);
assert (labels.dims.numDims == 2);
// Decoding the output tensor of MMYOLO to the NvDsInferParseObjectInfo format.
std::vector<NvDsInferParseObjectInfo> objects =
decodeMMYoloTensor(
(const int*)(num_dets.buffer),
(const float*)(bboxes.buffer),
(const float*)(scores.buffer),
(const int*)(labels.buffer),
conf_thres,
networkInfo.width,
networkInfo.height
);
objectList.clear();
objectList = objects;
return true;
}
/* Check that the custom function has been defined correctly */
CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomMMYOLO);
|