123 lines
4.2 KiB
C++
123 lines
4.2 KiB
C++
/*
|
|
* Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#include <algorithm>
|
|
#include <cassert>
|
|
#include <cmath>
|
|
#include <cstring>
|
|
#include <fstream>
|
|
#include <iostream>
|
|
#include <unordered_map>
|
|
#include "nvdsinfer_custom_impl.h"
|
|
|
|
#include <map>
|
|
|
|
#define kNMS_THRESH 0.45
|
|
|
|
static constexpr int LOCATIONS = 4;
|
|
struct alignas(float) Detection{
|
|
//center_x center_y w h
|
|
float bbox[LOCATIONS];
|
|
float conf; // bbox_conf * cls_conf
|
|
float class_id;
|
|
};
|
|
|
|
float iou(float lbox[4], float rbox[4]) {
|
|
float interBox[] = {
|
|
std::max(lbox[0] - lbox[2]/2.f , rbox[0] - rbox[2]/2.f), //left
|
|
std::min(lbox[0] + lbox[2]/2.f , rbox[0] + rbox[2]/2.f), //right
|
|
std::max(lbox[1] - lbox[3]/2.f , rbox[1] - rbox[3]/2.f), //top
|
|
std::min(lbox[1] + lbox[3]/2.f , rbox[1] + rbox[3]/2.f), //bottom
|
|
};
|
|
|
|
if(interBox[2] > interBox[3] || interBox[0] > interBox[1])
|
|
return 0.0f;
|
|
|
|
float interBoxS =(interBox[1]-interBox[0])*(interBox[3]-interBox[2]);
|
|
return interBoxS/(lbox[2]*lbox[3] + rbox[2]*rbox[3] -interBoxS);
|
|
}
|
|
|
|
bool cmp(Detection& a, Detection& b) {
|
|
return a.conf > b.conf;
|
|
}
|
|
|
|
void nms(std::vector<Detection>& res, float *output, float conf_thresh, float nms_thresh) {
|
|
int det_size = sizeof(Detection) / sizeof(float);
|
|
std::map<float, std::vector<Detection>> m;
|
|
for (int i = 0; i < output[0] && i < 1000; i++) {
|
|
if (output[1 + det_size * i + 4] <= conf_thresh) continue;
|
|
Detection det;
|
|
memcpy(&det, &output[1 + det_size * i], det_size * sizeof(float));
|
|
if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector<Detection>());
|
|
m[det.class_id].push_back(det);
|
|
}
|
|
for (auto it = m.begin(); it != m.end(); it++) {
|
|
auto& dets = it->second;
|
|
std::sort(dets.begin(), dets.end(), cmp);
|
|
for (size_t m = 0; m < dets.size(); ++m) {
|
|
auto& item = dets[m];
|
|
res.push_back(item);
|
|
for (size_t n = m + 1; n < dets.size(); ++n) {
|
|
if (iou(item.bbox, dets[n].bbox) > nms_thresh) {
|
|
dets.erase(dets.begin()+n);
|
|
--n;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/* This is a sample bounding box parsing function for the sample YoloV5 detector model */
|
|
static bool NvDsInferParseYoloV5(
|
|
std::vector<NvDsInferLayerInfo> const& outputLayersInfo,
|
|
NvDsInferNetworkInfo const& networkInfo,
|
|
NvDsInferParseDetectionParams const& detectionParams,
|
|
std::vector<NvDsInferParseObjectInfo>& objectList)
|
|
{
|
|
const float kCONF_THRESH = detectionParams.perClassThreshold[0];
|
|
|
|
std::vector<Detection> res;
|
|
|
|
nms(res, (float*)(outputLayersInfo[0].buffer), kCONF_THRESH, kNMS_THRESH);
|
|
|
|
for(auto& r : res) {
|
|
NvDsInferParseObjectInfo oinfo;
|
|
|
|
oinfo.classId = r.class_id;
|
|
oinfo.left = static_cast<unsigned int>(r.bbox[0]-r.bbox[2]*0.5f);
|
|
oinfo.top = static_cast<unsigned int>(r.bbox[1]-r.bbox[3]*0.5f);
|
|
oinfo.width = static_cast<unsigned int>(r.bbox[2]);
|
|
oinfo.height = static_cast<unsigned int>(r.bbox[3]);
|
|
oinfo.detectionConfidence = r.conf;
|
|
objectList.push_back(oinfo);
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
extern "C" bool NvDsInferParseCustomYoloV5(
|
|
std::vector<NvDsInferLayerInfo> const &outputLayersInfo,
|
|
NvDsInferNetworkInfo const &networkInfo,
|
|
NvDsInferParseDetectionParams const &detectionParams,
|
|
std::vector<NvDsInferParseObjectInfo> &objectList)
|
|
{
|
|
return NvDsInferParseYoloV5(
|
|
outputLayersInfo, networkInfo, detectionParams, objectList);
|
|
}
|
|
|
|
/* Check that the custom function has been defined correctly */
|
|
CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseCustomYoloV5);
|