-
Notifications
You must be signed in to change notification settings - Fork 0
/
people_counting.h
52 lines (40 loc) · 1.48 KB
/
people_counting.h
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
#ifndef _PC_H_
#define _PC_H_
#include <string>
#include <opencv2/opencv.hpp>
#include "NvInfer.h"
#include "yolo.h"
#include "jbuf.h"
#include "img.h"
#include "nvdsinfer_custom_impl.h"
const std::vector<float> kANCHORS = {
12, 23, 21, 41, 34, 54, 32, 112, 53,
78, 75, 120, 58, 196, 104, 288, 184, 322
};
const std::vector<std::vector<int>> kMASKS = {
{6, 7, 8},
{3, 4, 5},
{0, 1, 2}};
class PeopleDetector {
public:
static const int32_t input_tensor_height = 640;
static const int32_t input_tensor_width = 640;
static const int32_t input_tensor_depth = 3;
static constexpr const char* input_blob_name = "data";
std::vector<std::string> output_blob_names;
PeopleDetector(std::string cfg_path, std::string wts_path, int32_t batch_size, float cls_thres, float nms_thres, nvinfer1::ILogger& logger);
PeopleDetector(std::string model_path, int32_t batch_size, float cls_thres, float nms_thres, nvinfer1::ILogger& logger);
~PeopleDetector();
std::vector<NvDsInferParseObjectInfo> detect(cv::Mat img);
int32_t detect_capi(cv::Mat img, NvDsInferParseObjectInfo* boxes, int32_t& num);
private:
nvinfer1::ICudaEngine* init_engine(std::string cfg_path, std::string weight_path, nvinfer1::IBuilder* builder);
nvinfer1::ICudaEngine* engine;
nvinfer1::IExecutionContext* ctx;
UnifiedBufManager* buffers;
std::vector<NvDsInferLayerInfo> layer_info;
float nms_thres;
float cls_thres;
int32_t batch_size;
};
#endif