-
Notifications
You must be signed in to change notification settings - Fork 1
/
nms.h
71 lines (60 loc) · 1.89 KB
/
nms.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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
#ifndef NMS_H_
#define NMS_H_
template<typename T>
void nms(const T *pBoxes, const T *pScores, int nSample, double overlap,
int &nPick, T* &boxes_nms, T* &scores_nms) {
const int BOX_DIMS = 4;
vector<int> vPick(nSample);
vector<double> vArea(nSample);
for (int i = 0; i < nSample; ++i) {
vArea[i] = double(
pBoxes[i * BOX_DIMS + 2] - pBoxes[i * BOX_DIMS + 0] + 1)
* (pBoxes[i * BOX_DIMS + 3] - pBoxes[i * BOX_DIMS + 1] + 1);
if (vArea[i] < 0)
printf("Boxes area must >= 0");
}
std::multimap<T, int> scores;
for (int i = 0; i < nSample; ++i)
scores.insert(std::pair<T, int>(pScores[i], i));
nPick = 0;
do {
int last = scores.rbegin()->second;
vPick[nPick] = last;
nPick += 1;
for (typename std::multimap<T, int>::iterator it = scores.begin();
it != scores.end();) {
int it_idx = it->second;
T xx1 = max(pBoxes[0 + BOX_DIMS * last],
pBoxes[0 + BOX_DIMS * it_idx]);
T yy1 = max(pBoxes[1 + BOX_DIMS * last],
pBoxes[1 + BOX_DIMS * it_idx]);
T xx2 = min(pBoxes[2 + BOX_DIMS * last],
pBoxes[2 + BOX_DIMS * it_idx]);
T yy2 = min(pBoxes[3 + BOX_DIMS * last],
pBoxes[3 + BOX_DIMS * it_idx]);
double w = max(T(0.0), xx2 - xx1 + 1), h = max(T(0.0),
yy2 - yy1 + 1);
double ov = w * h / (vArea[last] + vArea[it_idx] - w * h);
if (ov > overlap) {
it = scores.erase(it);
} else {
it++;
}
}
} while (scores.size() != 0);
if (boxes_nms != NULL)
delete[] boxes_nms;
boxes_nms = new T[nPick * BOX_DIMS];
if (scores_nms != NULL)
delete[] scores_nms;
scores_nms = new T[nPick];
for (int i = 0; i < nPick; ++i) {
int index = vPick[i];
boxes_nms[i * BOX_DIMS + 0] = pBoxes[index * BOX_DIMS + 0];
boxes_nms[i * BOX_DIMS + 1] = pBoxes[index * BOX_DIMS + 1];
boxes_nms[i * BOX_DIMS + 2] = pBoxes[index * BOX_DIMS + 2];
boxes_nms[i * BOX_DIMS + 3] = pBoxes[index * BOX_DIMS + 3];
scores_nms[i] = pScores[index];
}
}
#endif