tnn_scrfd.h
4.1 KB
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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
//
// Created by DefTruth on 2021/12/30.
//
#ifndef LITE_AI_TOOLKIT_TNN_CV_TNN_SCRFD_H
#define LITE_AI_TOOLKIT_TNN_CV_TNN_SCRFD_H
#include "lite/tnn/core/tnn_core.h"
namespace tnncv
{
class LITE_EXPORTS TNNSCRFD : public BasicTNNHandler
{
public:
explicit TNNSCRFD(const std::string &_proto_path,
const std::string &_model_path,
unsigned int _num_threads = 1); //
~TNNSCRFD() override = default;
private:
// nested classes
typedef struct
{
float cx;
float cy;
float stride;
} SCRFDPoint;
typedef struct
{
float ratio;
int dw;
int dh;
bool flag;
} SCRFDScaleParams;
private:
// In TNN: x*scale + bias
std::vector<float> scale_vals = {1.f / 128.f, 1.f / 128.f, 1.f / 128.f}; // RGB
std::vector<float> bias_vals = {-127.5f / 128.f, -127.5f / 128.f, -127.5f / 128.f};
unsigned int fmc = 3; // feature map count
bool use_kps = false;
unsigned int num_anchors = 2;
std::vector<int> feat_stride_fpn = {8, 16, 32}; // steps, may [8, 16, 32, 64, 128]
// if num_anchors>1, then stack points in col major -> (height*num_anchor*width,2)
// anchor_centers = np.stack([anchor_centers]*self._num_anchors, axis=1).reshape( (-1,2) )
std::unordered_map<int, std::vector<SCRFDPoint>> center_points;
bool center_points_is_update = false;
static constexpr const unsigned int nms_pre = 1000;
static constexpr const unsigned int max_nms = 30000;
private:
void transform(const cv::Mat &mat_rs) override; // without resize
// initial steps and num_anchors
// https://github.com/deepinsight/insightface/blob/master/detection/scrfd/tools/scrfd.py
void initial_context();
void resize_unscale(const cv::Mat &mat,
cv::Mat &mat_rs,
int target_height,
int target_width,
SCRFDScaleParams &scale_params);
// generate once.
void generate_points(const int target_height, const int target_width);
void generate_bboxes_single_stride(const SCRFDScaleParams &scale_params,
std::shared_ptr<tnn::Mat> &score_pred,
std::shared_ptr<tnn::Mat> &bbox_pred,
unsigned int stride,
float score_threshold,
float img_height,
float img_width,
std::vector<types::BoxfWithLandmarks> &bbox_kps_collection);
void generate_bboxes_kps_single_stride(const SCRFDScaleParams &scale_params,
std::shared_ptr<tnn::Mat> &score_pred,
std::shared_ptr<tnn::Mat> &bbox_pred,
std::shared_ptr<tnn::Mat> &kps_pred,
unsigned int stride,
float score_threshold,
float img_height,
float img_width,
std::vector<types::BoxfWithLandmarks> &bbox_kps_collection);
void generate_bboxes_kps(const SCRFDScaleParams &scale_params,
std::vector<types::BoxfWithLandmarks> &bbox_kps_collection,
std::shared_ptr<tnn::Instance> &_instance,
float score_threshold, float img_height,
float img_width); // rescale & exclude
void nms_bboxes_kps(std::vector<types::BoxfWithLandmarks> &input,
std::vector<types::BoxfWithLandmarks> &output,
float iou_threshold, unsigned int topk);
public:
void detect(const cv::Mat &mat, std::vector<types::BoxfWithLandmarks> &detected_boxes_kps,
float score_threshold = 0.25f, float iou_threshold = 0.45f,
unsigned int topk = 400);
};
}
#endif //LITE_AI_TOOLKIT_TNN_CV_TNN_SCRFD_H