出处:nms.cc
#include <bits/stdc++.h>
namespace nms
{
struct proposal
{
float score, x1, y1, x2, y2;
};
inline static bool cmp(const proposal& a, const proposal& b)
{
return a.score < b.score;
}
inline static float iou(const proposal&, const proposal&) __attribute__((always_inline));
static float iou(const proposal& a, const proposal& b)
{
auto overlap = 0.f;
float iw = std::min(b.x2, a.x2) - std::max(b.x1, a.x1) + 1;
if (iw > 0) {
float ih = std::min(b.y2, a.y2) - std::max(b.y1, a.y1) + 1;
if (ih > 0) {
float ab = (b.x2 - b.x1 + 1) * (b.y2 - b.y1 + 1);
float aa = (a.x2 - a.x1 + 1) * (a.y2 - a.y1 + 1);
float inter = iw * ih;
overlap = inter / (aa + ab - inter);
}
}
return overlap;
}
enum class Method : uint32_t
{
LINEAR = 0,
GAUSSIAN,
HARD
};
size_t soft_nms(float* boxes,
int32_t* index,
size_t count,
Method method,
float Nt,
float sigma,
float threshold)
{
std::iota(index, index + count, 0); // np.arange()
auto p = reinterpret_cast<proposal*>(boxes);
auto N = count;
for (size_t i = 0; i < N; ++i) {
auto max = std::max_element(p + i, p + N, cmp);
std::swap(p[i], *max);
std::swap(index[i], index[max - p]);
auto j = i + 1;
auto weight = 0.f;
while (j < N) {
auto ov = iou(p[i], p[j]);
switch (method) {
case Method::LINEAR:
weight = ov > Nt ? 1.f - ov : 1.f;
break;
case Method::GAUSSIAN:
weight = std::exp(-(ov * ov) / sigma);
break;
case Method::HARD:
weight = ov > Nt ? 0.f : 1.f;
break;
}
p[j].score *= weight;
if (p[j].score < threshold) {
N--;
std::swap(p[j], p[N]);
std::swap(index[j], index[N]);
j--;
}
j++;
}
};
return N;
}
} /* namespace nms */