Spaces:
Build error
Build error
/** | |
* 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); | |