158 lines
4.9 KiB
C++
158 lines
4.9 KiB
C++
/**
|
|
* SCRFD Detector Implementation
|
|
*/
|
|
|
|
#include "face/scrfd_detector.h"
|
|
#include "ai_scheduler.h" // For BorrowedOutput
|
|
#include "face/face_detection_utils.h"
|
|
#include <algorithm>
|
|
#include <cstring>
|
|
|
|
namespace rk3588 {
|
|
|
|
ScrfdDetector::ScrfdDetector() = default;
|
|
ScrfdDetector::~ScrfdDetector() = default;
|
|
|
|
void ScrfdDetector::Init(int model_w, int model_h) {
|
|
model_w_ = model_w;
|
|
model_h_ = model_h;
|
|
|
|
// Generate center points
|
|
const int strides[] = {8, 16, 32};
|
|
|
|
for (int stride : strides) {
|
|
int num_grid = model_w_ / stride;
|
|
for (int y = 0; y < num_grid; ++y) {
|
|
for (int x = 0; x < num_grid; ++x) {
|
|
// 2 anchors per location
|
|
for (int a = 0; a < 2; ++a) {
|
|
CenterPoint pt;
|
|
pt.cx = static_cast<float>(x);
|
|
pt.cy = static_cast<float>(y);
|
|
pt.stride = static_cast<float>(stride);
|
|
center_points_.push_back(pt);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
std::vector<FaceDetItem> ScrfdDetector::Decode(
|
|
const std::vector<AiScheduler::BorrowedOutput>& outputs,
|
|
int src_w, int src_h,
|
|
const ScrfdConfig& cfg) {
|
|
|
|
std::vector<FaceDetItem> detections;
|
|
|
|
if (outputs.size() != 9) return detections;
|
|
|
|
// Output order: score_8, score_16, score_32, bbox_8, bbox_16, bbox_32, kps_8, kps_16, kps_32
|
|
const int anchor_counts[] = {12800, 3200, 800};
|
|
const int strides[] = {8, 16, 32};
|
|
|
|
size_t anchor_idx = 0;
|
|
|
|
for (int s = 0; s < 3; ++s) {
|
|
int stride = strides[s];
|
|
int count = anchor_counts[s];
|
|
|
|
const auto& score_out = outputs[s];
|
|
const auto& bbox_out = outputs[s + 3];
|
|
const auto& kps_out = outputs[s + 6];
|
|
|
|
if (score_out.dims.size() < 3) continue;
|
|
|
|
const float* scores = reinterpret_cast<const float*>(score_out.data);
|
|
const float* bboxes = reinterpret_cast<const float*>(bbox_out.data);
|
|
const float* kps = reinterpret_cast<const float*>(kps_out.data);
|
|
|
|
if (!scores || !bboxes || !kps) continue;
|
|
|
|
for (int i = 0; i < count; ++i) {
|
|
if (anchor_idx >= center_points_.size()) break;
|
|
|
|
float score = scores[i];
|
|
if (score < cfg.conf_thresh) {
|
|
anchor_idx++;
|
|
continue;
|
|
}
|
|
|
|
const CenterPoint& pt = center_points_[anchor_idx];
|
|
|
|
// BBox: [left, top, right, bottom] - distances from center
|
|
float left = bboxes[i * 4 + 0];
|
|
float top = bboxes[i * 4 + 1];
|
|
float right = bboxes[i * 4 + 2];
|
|
float bottom = bboxes[i * 4 + 3];
|
|
|
|
// Decode to image coordinates (640x640)
|
|
float x1_640 = (pt.cx - left) * stride;
|
|
float y1_640 = (pt.cy - top) * stride;
|
|
float x2_640 = (pt.cx + right) * stride;
|
|
float y2_640 = (pt.cy + bottom) * stride;
|
|
|
|
// Scale to original image size
|
|
float scale_x = static_cast<float>(src_w) / model_w_;
|
|
float scale_y = static_cast<float>(src_h) / model_h_;
|
|
|
|
FaceDetItem det;
|
|
det.bbox.x = x1_640 * scale_x;
|
|
det.bbox.y = y1_640 * scale_y;
|
|
det.bbox.w = (x2_640 - x1_640) * scale_x;
|
|
det.bbox.h = (y2_640 - y1_640) * scale_y;
|
|
det.score = score;
|
|
det.has_landmarks = cfg.output_landmarks;
|
|
|
|
// Keypoints
|
|
if (cfg.output_landmarks) {
|
|
for (int p = 0; p < 5; ++p) {
|
|
float kps_x = kps[i * 10 + p * 2 + 0];
|
|
float kps_y = kps[i * 10 + p * 2 + 1];
|
|
float kx_640 = (pt.cx + kps_x) * stride;
|
|
float ky_640 = (pt.cy + kps_y) * stride;
|
|
det.landmarks[p].x = kx_640 * scale_x;
|
|
det.landmarks[p].y = ky_640 * scale_y;
|
|
}
|
|
}
|
|
|
|
detections.push_back(det);
|
|
anchor_idx++;
|
|
}
|
|
}
|
|
|
|
return detections;
|
|
}
|
|
|
|
std::vector<FaceDetItem> ScrfdDetector::ApplyNMS(
|
|
std::vector<FaceDetItem>& dets,
|
|
float nms_thresh) {
|
|
|
|
if (dets.empty()) return dets;
|
|
|
|
// Sort by score
|
|
std::sort(dets.begin(), dets.end(),
|
|
[](const FaceDetItem& a, const FaceDetItem& b) {
|
|
return a.score > b.score;
|
|
});
|
|
|
|
std::vector<FaceDetItem> keep;
|
|
std::vector<bool> suppressed(dets.size(), false);
|
|
|
|
for (size_t i = 0; i < dets.size(); ++i) {
|
|
if (suppressed[i]) continue;
|
|
|
|
keep.push_back(dets[i]);
|
|
|
|
for (size_t j = i + 1; j < dets.size(); ++j) {
|
|
if (suppressed[j]) continue;
|
|
if (face_detection::IoU(dets[i].bbox, dets[j].bbox) > nms_thresh) {
|
|
suppressed[j] = true;
|
|
}
|
|
}
|
|
}
|
|
|
|
return keep;
|
|
}
|
|
|
|
} // namespace rk3588
|