485 lines
17 KiB
C++
485 lines
17 KiB
C++
|
//
|
||
|
|
// Created by DefTruth on 2021/10/7.
|
||
|
|
//
|
||
|
|
|
||
|
|
#include "utils.h"
|
||
|
|
|
||
|
|
//*************************************** lite::utils **********************************************//
|
||
|
|
std::string lite::utils::to_string(const std::wstring &wstr)
|
||
|
|
{
|
||
|
|
unsigned len = wstr.size() * 4;
|
||
|
|
setlocale(LC_CTYPE, "");
|
||
|
|
char *p = new char[len];
|
||
|
|
wcstombs(p, wstr.c_str(), len);
|
||
|
|
std::string str(p);
|
||
|
|
delete[] p;
|
||
|
|
return str;
|
||
|
|
}
|
||
|
|
|
||
|
|
std::wstring lite::utils::to_wstring(const std::string &str)
|
||
|
|
{
|
||
|
|
unsigned len = str.size() * 2;
|
||
|
|
setlocale(LC_CTYPE, "");
|
||
|
|
wchar_t *p = new wchar_t[len];
|
||
|
|
mbstowcs(p, str.c_str(), len);
|
||
|
|
std::wstring wstr(p);
|
||
|
|
delete[] p;
|
||
|
|
return wstr;
|
||
|
|
}
|
||
|
|
|
||
|
|
// reference: https://github.com/DefTruth/headpose-fsanet-pytorch/blob/master/src/utils.py
|
||
|
|
void lite::utils::draw_axis_inplace(cv::Mat &mat_inplace,
|
||
|
|
const types::EulerAngles &euler_angles,
|
||
|
|
float size, int thickness)
|
||
|
|
{
|
||
|
|
if (!euler_angles.flag) return;
|
||
|
|
|
||
|
|
const float pitch = euler_angles.pitch * _PI / 180.f;
|
||
|
|
const float yaw = -euler_angles.yaw * _PI / 180.f;
|
||
|
|
const float roll = euler_angles.roll * _PI / 180.f;
|
||
|
|
|
||
|
|
const float height = static_cast<float>(mat_inplace.rows);
|
||
|
|
const float width = static_cast<float>(mat_inplace.cols);
|
||
|
|
|
||
|
|
const int tdx = static_cast<int>(width / 2.0f);
|
||
|
|
const int tdy = static_cast<int>(height / 2.0f);
|
||
|
|
|
||
|
|
// X-Axis pointing to right. drawn in red
|
||
|
|
const int x1 = static_cast<int>(size * std::cos(yaw) * std::cos(roll)) + tdx;
|
||
|
|
const int y1 = static_cast<int>(
|
||
|
|
size * (std::cos(pitch) * std::sin(roll)
|
||
|
|
+ std::cos(roll) * std::sin(pitch) * std::sin(yaw))
|
||
|
|
) + tdy;
|
||
|
|
// Y-Axis | drawn in green
|
||
|
|
const int x2 = static_cast<int>(-size * std::cos(yaw) * std::sin(roll)) + tdx;
|
||
|
|
const int y2 = static_cast<int>(
|
||
|
|
size * (std::cos(pitch) * std::cos(roll)
|
||
|
|
- std::sin(pitch) * std::sin(yaw) * std::sin(roll))
|
||
|
|
) + tdy;
|
||
|
|
// Z-Axis (out of the screen) drawn in blue
|
||
|
|
const int x3 = static_cast<int>(size * std::sin(yaw)) + tdx;
|
||
|
|
const int y3 = static_cast<int>(-size * std::cos(yaw) * std::sin(pitch)) + tdy;
|
||
|
|
|
||
|
|
cv::line(mat_inplace, cv::Point2i(tdx, tdy), cv::Point2i(x1, y1), cv::Scalar(0, 0, 255), thickness);
|
||
|
|
cv::line(mat_inplace, cv::Point2i(tdx, tdy), cv::Point2i(x2, y2), cv::Scalar(0, 255, 0), thickness);
|
||
|
|
cv::line(mat_inplace, cv::Point2i(tdx, tdy), cv::Point2i(x3, y3), cv::Scalar(255, 0, 0), thickness);
|
||
|
|
}
|
||
|
|
|
||
|
|
cv::Mat lite::utils::draw_axis(const cv::Mat &mat,
|
||
|
|
const types::EulerAngles &euler_angles,
|
||
|
|
float size, int thickness)
|
||
|
|
{
|
||
|
|
if (!euler_angles.flag) return mat;
|
||
|
|
|
||
|
|
cv::Mat mat_copy = mat.clone();
|
||
|
|
const float pitch = euler_angles.pitch * _PI / 180.f;
|
||
|
|
const float yaw = -euler_angles.yaw * _PI / 180.f;
|
||
|
|
const float roll = euler_angles.roll * _PI / 180.f;
|
||
|
|
|
||
|
|
const float height = static_cast<float>(mat_copy.rows);
|
||
|
|
const float width = static_cast<float>(mat_copy.cols);
|
||
|
|
|
||
|
|
const int tdx = static_cast<int>(width / 2.0f);
|
||
|
|
const int tdy = static_cast<int>(height / 2.0f);
|
||
|
|
|
||
|
|
// X-Axis pointing to right. drawn in red
|
||
|
|
const int x1 = static_cast<int>(size * std::cos(yaw) * std::cos(roll)) + tdx;
|
||
|
|
const int y1 = static_cast<int>(
|
||
|
|
size * (std::cos(pitch) * std::sin(roll)
|
||
|
|
+ std::cos(roll) * std::sin(pitch) * std::sin(yaw))
|
||
|
|
) + tdy;
|
||
|
|
// Y-Axis | drawn in green
|
||
|
|
const int x2 = static_cast<int>(-size * std::cos(yaw) * std::sin(roll)) + tdx;
|
||
|
|
const int y2 = static_cast<int>(
|
||
|
|
size * (std::cos(pitch) * std::cos(roll)
|
||
|
|
- std::sin(pitch) * std::sin(yaw) * std::sin(roll))
|
||
|
|
) + tdy;
|
||
|
|
// Z-Axis (out of the screen) drawn in blue
|
||
|
|
const int x3 = static_cast<int>(size * std::sin(yaw)) + tdx;
|
||
|
|
const int y3 = static_cast<int>(-size * std::cos(yaw) * std::sin(pitch)) + tdy;
|
||
|
|
|
||
|
|
cv::line(mat_copy, cv::Point2i(tdx, tdy), cv::Point2i(x1, y1), cv::Scalar(0, 0, 255), thickness);
|
||
|
|
cv::line(mat_copy, cv::Point2i(tdx, tdy), cv::Point2i(x2, y2), cv::Scalar(0, 255, 0), thickness);
|
||
|
|
cv::line(mat_copy, cv::Point2i(tdx, tdy), cv::Point2i(x3, y3), cv::Scalar(255, 0, 0), thickness);
|
||
|
|
|
||
|
|
return mat_copy;
|
||
|
|
}
|
||
|
|
|
||
|
|
cv::Mat lite::utils::draw_landmarks(const cv::Mat &mat, types::Landmarks &landmarks)
|
||
|
|
{
|
||
|
|
if (landmarks.points.empty() || !landmarks.flag) return mat;
|
||
|
|
cv::Mat mat_copy = mat.clone();
|
||
|
|
for (const auto &point: landmarks.points)
|
||
|
|
cv::circle(mat_copy, point, 2, cv::Scalar(0, 255, 0), -1);
|
||
|
|
return mat_copy;
|
||
|
|
}
|
||
|
|
|
||
|
|
void lite::utils::draw_landmarks_inplace(cv::Mat &mat, types::Landmarks &landmarks)
|
||
|
|
{
|
||
|
|
if (landmarks.points.empty() || !landmarks.flag) return;
|
||
|
|
for (const auto &point: landmarks.points)
|
||
|
|
cv::circle(mat, point, 2, cv::Scalar(0, 255, 0), -1);
|
||
|
|
}
|
||
|
|
|
||
|
|
void lite::utils::draw_boxes_inplace(cv::Mat &mat_inplace, const std::vector<types::Boxf> &boxes)
|
||
|
|
{
|
||
|
|
if (boxes.empty()) return;
|
||
|
|
for (const auto &box: boxes)
|
||
|
|
{
|
||
|
|
if (box.flag)
|
||
|
|
{
|
||
|
|
cv::rectangle(mat_inplace, box.rect(), cv::Scalar(255, 255, 0), 2);
|
||
|
|
if (box.label_text)
|
||
|
|
{
|
||
|
|
std::string label_text(box.label_text);
|
||
|
|
label_text = label_text + ":" + std::to_string(box.score).substr(0, 4);
|
||
|
|
cv::putText(mat_inplace, label_text, box.tl(), cv::FONT_HERSHEY_SIMPLEX,
|
||
|
|
0.6f, cv::Scalar(0, 255, 0), 2);
|
||
|
|
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
cv::Mat lite::utils::draw_boxes(const cv::Mat &mat, const std::vector<types::Boxf> &boxes)
|
||
|
|
{
|
||
|
|
if (boxes.empty()) return mat;
|
||
|
|
cv::Mat canva = mat.clone();
|
||
|
|
for (const auto &box: boxes)
|
||
|
|
{
|
||
|
|
if (box.flag)
|
||
|
|
{
|
||
|
|
cv::rectangle(canva, box.rect(), cv::Scalar(255, 255, 0), 2);
|
||
|
|
if (box.label_text)
|
||
|
|
{
|
||
|
|
std::string label_text(box.label_text);
|
||
|
|
label_text = label_text + ":" + std::to_string(box.score).substr(0, 4);
|
||
|
|
cv::putText(canva, label_text, box.tl(), cv::FONT_HERSHEY_SIMPLEX,
|
||
|
|
0.6f, cv::Scalar(0, 255, 0), 2);
|
||
|
|
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return canva;
|
||
|
|
}
|
||
|
|
|
||
|
|
void lite::utils::draw_boxes_with_landmarks_inplace(cv::Mat &mat_inplace, const std::vector<types::BoxfWithLandmarks> &boxes_kps, bool text)
|
||
|
|
{
|
||
|
|
if (boxes_kps.empty()) return;
|
||
|
|
for (const auto &box_kps: boxes_kps)
|
||
|
|
{
|
||
|
|
if (box_kps.flag)
|
||
|
|
{
|
||
|
|
// box
|
||
|
|
if (box_kps.box.flag)
|
||
|
|
{
|
||
|
|
cv::rectangle(mat_inplace, box_kps.box.rect(), cv::Scalar(255, 255, 0), 2);
|
||
|
|
if (box_kps.box.label_text && text)
|
||
|
|
{
|
||
|
|
std::string label_text(box_kps.box.label_text);
|
||
|
|
label_text = label_text + ":" + std::to_string(box_kps.box.score).substr(0, 4);
|
||
|
|
cv::putText(mat_inplace, label_text, box_kps.box.tl(), cv::FONT_HERSHEY_SIMPLEX,
|
||
|
|
0.6f, cv::Scalar(0, 255, 0), 2);
|
||
|
|
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// landmarks
|
||
|
|
if (box_kps.landmarks.flag && !box_kps.landmarks.points.empty())
|
||
|
|
{
|
||
|
|
for (const auto &point: box_kps.landmarks.points)
|
||
|
|
cv::circle(mat_inplace, point, 2, cv::Scalar(0, 255, 0), -1);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
cv::Mat lite::utils::draw_boxes_with_landmarks(const cv::Mat &mat, const std::vector<types::BoxfWithLandmarks> &boxes_kps, bool text)
|
||
|
|
{
|
||
|
|
if (boxes_kps.empty()) return mat;
|
||
|
|
cv::Mat canva = mat.clone();
|
||
|
|
for (const auto &box_kps: boxes_kps)
|
||
|
|
{
|
||
|
|
if (box_kps.flag)
|
||
|
|
{
|
||
|
|
// box
|
||
|
|
if (box_kps.box.flag)
|
||
|
|
{
|
||
|
|
cv::rectangle(canva, box_kps.box.rect(), cv::Scalar(255, 255, 0), 2);
|
||
|
|
if (box_kps.box.label_text && text)
|
||
|
|
{
|
||
|
|
std::string label_text(box_kps.box.label_text);
|
||
|
|
label_text = label_text + ":" + std::to_string(box_kps.box.score).substr(0, 4);
|
||
|
|
cv::putText(canva, label_text, box_kps.box.tl(), cv::FONT_HERSHEY_SIMPLEX,
|
||
|
|
0.6f, cv::Scalar(0, 255, 0), 2);
|
||
|
|
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// landmarks
|
||
|
|
if (box_kps.landmarks.flag && !box_kps.landmarks.points.empty())
|
||
|
|
{
|
||
|
|
for (const auto &point: box_kps.landmarks.points)
|
||
|
|
cv::circle(canva, point, 2, cv::Scalar(0, 255, 0), -1);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return canva;
|
||
|
|
}
|
||
|
|
|
||
|
|
cv::Mat lite::utils::draw_age(const cv::Mat &mat, types::Age &age)
|
||
|
|
{
|
||
|
|
if (!age.flag) return mat;
|
||
|
|
cv::Mat canva = mat.clone();
|
||
|
|
const unsigned int offset = static_cast<unsigned int>(
|
||
|
|
0.1f * static_cast<float>(mat.rows));
|
||
|
|
std::string age_text = "Age:" + std::to_string(age.age).substr(0, 4);
|
||
|
|
std::string interval_text = "Interval" + std::to_string(age.age_interval[0])
|
||
|
|
+ "_" + std::to_string(age.age_interval[1]);
|
||
|
|
std::string interval_prob = "Prob:" + std::to_string(age.interval_prob).substr(0, 4);
|
||
|
|
cv::putText(canva, age_text, cv::Point2i(10, offset),
|
||
|
|
cv::FONT_HERSHEY_SIMPLEX, 0.6f, cv::Scalar(0, 255, 0), 2);
|
||
|
|
cv::putText(canva, interval_text, cv::Point2i(10, offset * 2),
|
||
|
|
cv::FONT_HERSHEY_SIMPLEX, 0.6f, cv::Scalar(255, 0, 0), 2);
|
||
|
|
cv::putText(canva, interval_prob, cv::Point2i(10, offset * 3),
|
||
|
|
cv::FONT_HERSHEY_SIMPLEX, 0.6f, cv::Scalar(0, 0, 255), 2);
|
||
|
|
return canva;
|
||
|
|
}
|
||
|
|
|
||
|
|
void lite::utils::draw_age_inplace(cv::Mat &mat_inplace, types::Age &age)
|
||
|
|
{
|
||
|
|
if (!age.flag) return;
|
||
|
|
const unsigned int offset = static_cast<unsigned int>(
|
||
|
|
0.1f * static_cast<float>(mat_inplace.rows));
|
||
|
|
std::string age_text = "Age:" + std::to_string(age.age).substr(0, 4);
|
||
|
|
std::string interval_text = "Interval" + std::to_string(age.age_interval[0])
|
||
|
|
+ "_" + std::to_string(age.age_interval[1]);
|
||
|
|
std::string interval_prob = "Prob:" + std::to_string(age.interval_prob).substr(0, 4);
|
||
|
|
cv::putText(mat_inplace, age_text, cv::Point2i(10, offset),
|
||
|
|
cv::FONT_HERSHEY_SIMPLEX, 0.6f, cv::Scalar(0, 255, 0), 2);
|
||
|
|
cv::putText(mat_inplace, interval_text, cv::Point2i(10, offset * 2),
|
||
|
|
cv::FONT_HERSHEY_SIMPLEX, 0.6f, cv::Scalar(255, 0, 0), 2);
|
||
|
|
cv::putText(mat_inplace, interval_prob, cv::Point2i(10, offset * 3),
|
||
|
|
cv::FONT_HERSHEY_SIMPLEX, 0.6f, cv::Scalar(0, 0, 255), 2);
|
||
|
|
}
|
||
|
|
|
||
|
|
cv::Mat lite::utils::draw_gender(const cv::Mat &mat, types::Gender &gender)
|
||
|
|
{
|
||
|
|
if (!gender.flag) return mat;
|
||
|
|
cv::Mat canva = mat.clone();
|
||
|
|
const unsigned int offset = static_cast<unsigned int>(
|
||
|
|
0.1f * static_cast<float>(mat.rows));
|
||
|
|
std::string gender_text = std::to_string(gender.label) + ":"
|
||
|
|
+ gender.text + ":" + std::to_string(gender.score).substr(0, 4);
|
||
|
|
cv::putText(canva, gender_text, cv::Point2i(10, offset),
|
||
|
|
cv::FONT_HERSHEY_SIMPLEX, 0.6f, cv::Scalar(0, 255, 0), 2);
|
||
|
|
return canva;
|
||
|
|
}
|
||
|
|
|
||
|
|
void lite::utils::draw_gender_inplace(cv::Mat &mat_inplace, types::Gender &gender)
|
||
|
|
{
|
||
|
|
if (!gender.flag) return;
|
||
|
|
const unsigned int offset = static_cast<unsigned int>(
|
||
|
|
0.1f * static_cast<float>(mat_inplace.rows));
|
||
|
|
std::string gender_text = std::to_string(gender.label) + ":"
|
||
|
|
+ gender.text + ":" + std::to_string(gender.score).substr(0, 4);
|
||
|
|
cv::putText(mat_inplace, gender_text, cv::Point2i(10, offset),
|
||
|
|
cv::FONT_HERSHEY_SIMPLEX, 0.6f, cv::Scalar(0, 255, 0), 2);
|
||
|
|
}
|
||
|
|
|
||
|
|
cv::Mat lite::utils::draw_emotion(const cv::Mat &mat, types::Emotions &emotions)
|
||
|
|
{
|
||
|
|
if (!emotions.flag) return mat;
|
||
|
|
cv::Mat canva = mat.clone();
|
||
|
|
const unsigned int offset = static_cast<unsigned int>(
|
||
|
|
0.1f * static_cast<float>(mat.rows));
|
||
|
|
std::string emotion_text = std::to_string(emotions.label) + ":"
|
||
|
|
+ emotions.text + ":" + std::to_string(emotions.score).substr(0, 4);
|
||
|
|
cv::putText(canva, emotion_text, cv::Point2i(10, offset),
|
||
|
|
cv::FONT_HERSHEY_SIMPLEX, 0.6f, cv::Scalar(0, 255, 0), 2);
|
||
|
|
return canva;
|
||
|
|
}
|
||
|
|
|
||
|
|
void lite::utils::draw_emotion_inplace(cv::Mat &mat_inplace, types::Emotions &emotions)
|
||
|
|
{
|
||
|
|
if (!emotions.flag) return;
|
||
|
|
const unsigned int offset = static_cast<unsigned int>(
|
||
|
|
0.1f * static_cast<float>(mat_inplace.rows));
|
||
|
|
std::string emotion_text = std::to_string(emotions.label) + ":"
|
||
|
|
+ emotions.text + ":" + std::to_string(emotions.score).substr(0, 4);
|
||
|
|
cv::putText(mat_inplace, emotion_text, cv::Point2i(10, offset),
|
||
|
|
cv::FONT_HERSHEY_SIMPLEX, 0.6f, cv::Scalar(0, 255, 0), 2);
|
||
|
|
}
|
||
|
|
|
||
|
|
// reference: https://github.com/Linzaer/Ultra-Light-Fast-Generic-Face-Detector-1MB/
|
||
|
|
// blob/master/ncnn/src/UltraFace.cpp
|
||
|
|
void lite::utils::hard_nms(std::vector<types::Boxf> &input, std::vector<types::Boxf> &output,
|
||
|
|
float iou_threshold, unsigned int topk)
|
||
|
|
{
|
||
|
|
if (input.empty()) return;
|
||
|
|
std::sort(input.begin(), input.end(),
|
||
|
|
[](const types::Boxf &a, const types::Boxf &b)
|
||
|
|
{ return a.score > b.score; });
|
||
|
|
const unsigned int box_num = input.size();
|
||
|
|
std::vector<int> merged(box_num, 0);
|
||
|
|
|
||
|
|
unsigned int count = 0;
|
||
|
|
for (unsigned int i = 0; i < box_num; ++i)
|
||
|
|
{
|
||
|
|
if (merged[i]) continue;
|
||
|
|
std::vector<types::Boxf> buf;
|
||
|
|
|
||
|
|
buf.push_back(input[i]);
|
||
|
|
merged[i] = 1;
|
||
|
|
|
||
|
|
for (unsigned int j = i + 1; j < box_num; ++j)
|
||
|
|
{
|
||
|
|
if (merged[j]) continue;
|
||
|
|
|
||
|
|
float iou = static_cast<float>(input[i].iou_of(input[j]));
|
||
|
|
|
||
|
|
if (iou > iou_threshold)
|
||
|
|
{
|
||
|
|
merged[j] = 1;
|
||
|
|
buf.push_back(input[j]);
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
output.push_back(buf[0]);
|
||
|
|
|
||
|
|
// keep top k
|
||
|
|
count += 1;
|
||
|
|
if (count >= topk)
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
void lite::utils::blending_nms(std::vector<types::Boxf> &input, std::vector<types::Boxf> &output,
|
||
|
|
float iou_threshold, unsigned int topk)
|
||
|
|
{
|
||
|
|
if (input.empty()) return;
|
||
|
|
std::sort(input.begin(), input.end(),
|
||
|
|
[](const types::Boxf &a, const types::Boxf &b)
|
||
|
|
{ return a.score > b.score; });
|
||
|
|
const unsigned int box_num = input.size();
|
||
|
|
std::vector<int> merged(box_num, 0);
|
||
|
|
|
||
|
|
unsigned int count = 0;
|
||
|
|
for (unsigned int i = 0; i < box_num; ++i)
|
||
|
|
{
|
||
|
|
if (merged[i]) continue;
|
||
|
|
std::vector<types::Boxf> buf;
|
||
|
|
|
||
|
|
buf.push_back(input[i]);
|
||
|
|
merged[i] = 1;
|
||
|
|
|
||
|
|
for (unsigned int j = i + 1; j < box_num; ++j)
|
||
|
|
{
|
||
|
|
if (merged[j]) continue;
|
||
|
|
|
||
|
|
float iou = static_cast<float>(input[i].iou_of(input[j]));
|
||
|
|
if (iou > iou_threshold)
|
||
|
|
{
|
||
|
|
merged[j] = 1;
|
||
|
|
buf.push_back(input[j]);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
float total = 0.f;
|
||
|
|
for (unsigned int k = 0; k < buf.size(); ++k)
|
||
|
|
{
|
||
|
|
total += std::exp(buf[k].score);
|
||
|
|
}
|
||
|
|
types::Boxf rects;
|
||
|
|
for (unsigned int l = 0; l < buf.size(); ++l)
|
||
|
|
{
|
||
|
|
float rate = std::exp(buf[l].score) / total;
|
||
|
|
rects.x1 += buf[l].x1 * rate;
|
||
|
|
rects.y1 += buf[l].y1 * rate;
|
||
|
|
rects.x2 += buf[l].x2 * rate;
|
||
|
|
rects.y2 += buf[l].y2 * rate;
|
||
|
|
rects.score += buf[l].score * rate;
|
||
|
|
}
|
||
|
|
rects.flag = true;
|
||
|
|
output.push_back(rects);
|
||
|
|
|
||
|
|
// keep top k
|
||
|
|
count += 1;
|
||
|
|
if (count >= topk)
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// reference: https://github.com/ultralytics/yolov5/blob/master/utils/general.py
|
||
|
|
void lite::utils::offset_nms(std::vector<types::Boxf> &input, std::vector<types::Boxf> &output,
|
||
|
|
float iou_threshold, unsigned int topk)
|
||
|
|
{
|
||
|
|
if (input.empty()) return;
|
||
|
|
std::sort(input.begin(), input.end(),
|
||
|
|
[](const types::Boxf &a, const types::Boxf &b)
|
||
|
|
{ return a.score > b.score; });
|
||
|
|
const unsigned int box_num = input.size();
|
||
|
|
std::vector<int> merged(box_num, 0);
|
||
|
|
|
||
|
|
const float offset = 4096.f;
|
||
|
|
/** Add offset according to classes.
|
||
|
|
* That is, separate the boxes into categories, and each category performs its
|
||
|
|
* own NMS operation. The same offset will be used for those predicted to be of
|
||
|
|
* the same category. Therefore, the relative positions of boxes of the same
|
||
|
|
* category will remain unchanged. Box of different classes will be farther away
|
||
|
|
* after offset, because offsets are different. In this way, some overlapping but
|
||
|
|
* different categories of entities are not filtered out by the NMS. Very clever!
|
||
|
|
*/
|
||
|
|
for (unsigned int i = 0; i < box_num; ++i)
|
||
|
|
{
|
||
|
|
input[i].x1 += static_cast<float>(input[i].label) * offset;
|
||
|
|
input[i].y1 += static_cast<float>(input[i].label) * offset;
|
||
|
|
input[i].x2 += static_cast<float>(input[i].label) * offset;
|
||
|
|
input[i].y2 += static_cast<float>(input[i].label) * offset;
|
||
|
|
}
|
||
|
|
|
||
|
|
unsigned int count = 0;
|
||
|
|
for (unsigned int i = 0; i < box_num; ++i)
|
||
|
|
{
|
||
|
|
if (merged[i]) continue;
|
||
|
|
std::vector<types::Boxf> buf;
|
||
|
|
|
||
|
|
buf.push_back(input[i]);
|
||
|
|
merged[i] = 1;
|
||
|
|
|
||
|
|
for (unsigned int j = i + 1; j < box_num; ++j)
|
||
|
|
{
|
||
|
|
if (merged[j]) continue;
|
||
|
|
|
||
|
|
float iou = static_cast<float>(input[i].iou_of(input[j]));
|
||
|
|
|
||
|
|
if (iou > iou_threshold)
|
||
|
|
{
|
||
|
|
merged[j] = 1;
|
||
|
|
buf.push_back(input[j]);
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
output.push_back(buf[0]);
|
||
|
|
|
||
|
|
// keep top k
|
||
|
|
count += 1;
|
||
|
|
if (count >= topk)
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
|
||
|
|
/** Substract offset.*/
|
||
|
|
if (!output.empty())
|
||
|
|
{
|
||
|
|
for (unsigned int i = 0; i < output.size(); ++i)
|
||
|
|
{
|
||
|
|
output[i].x1 -= static_cast<float>(output[i].label) * offset;
|
||
|
|
output[i].y1 -= static_cast<float>(output[i].label) * offset;
|
||
|
|
output[i].x2 -= static_cast<float>(output[i].label) * offset;
|
||
|
|
output[i].y2 -= static_cast<float>(output[i].label) * offset;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
//*************************************** lite::utils **********************************************//
|