dtourolle
d753062c6c
Source (KPN++ pipeline nodes, ArcFace embedders, SCRFD/YuNet detectors, gallery builder), build scripts, and eval artifacts. - external/KPN as a git submodule (gitea.tourolle.paris/dtourolle/KPN) - ONNX models tracked via Git LFS (models/*.onnx) - generated outputs, TensorRT engines, reference repos, and media ignored
231 lines
10 KiB
C++
231 lines
10 KiB
C++
#pragma once
|
||
#include "ort_provider.hpp"
|
||
#include "types.hpp"
|
||
|
||
#include <onnxruntime/onnxruntime_cxx_api.h>
|
||
#include <opencv2/dnn.hpp>
|
||
#include <opencv2/imgproc.hpp>
|
||
|
||
#include <memory>
|
||
#include <stdexcept>
|
||
#include <string>
|
||
#include <vector>
|
||
|
||
// ── SCRFDDecoder ──────────────────────────────────────────────────────────────
|
||
// ONNX Runtime-based decoder for InsightFace SCRFD face detector with kps.
|
||
// Uses ORT instead of cv::dnn because OpenCV 4.x cannot load SCRFD's dynamic
|
||
// Shape nodes. ORT handles dynamic shapes natively and is thread-safe for
|
||
// concurrent Run() calls.
|
||
//
|
||
// Model output layout (9 tensors, InsightFace export order):
|
||
// [0-2] score_s8 / score_s16 / score_s32 — flat (N,)
|
||
// [3-5] bbox_s8 / bbox_s16 / bbox_s32 — flat (N*4,) distance format
|
||
// [6-8] kps_s8 / kps_s16 / kps_s32 — flat (N*10,) distance format
|
||
//
|
||
// Landmark order (same as YuNet/ArcFace convention):
|
||
// [0] right-eye [1] left-eye [2] nose [3] right-mouth [4] left-mouth
|
||
|
||
struct SCRFDDecoder {
|
||
static constexpr int kInputW = 640;
|
||
static constexpr int kInputH = 640;
|
||
static constexpr int kAllStrides[4] = {8, 16, 32, 64};
|
||
static constexpr int kAnchors = 2;
|
||
|
||
SCRFDDecoder(const std::string& model_path,
|
||
float conf_threshold, float nms_threshold,
|
||
OrtProvider provider = OrtProvider::CPU,
|
||
TrtConfig trt_cfg = {})
|
||
: conf_threshold_(conf_threshold)
|
||
, nms_threshold_(nms_threshold)
|
||
{
|
||
Ort::SessionOptions opts;
|
||
opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
||
opts.SetIntraOpNumThreads(1);
|
||
|
||
// SCRFD ONNX has a dynamic H/W input; we letterbox to 640×640 at
|
||
// runtime, so pin the TRT profile to that single shape — otherwise
|
||
// TRT picks generic shapes and either rebuilds per-call or falls
|
||
// back to CUDA EP.
|
||
if (provider == OrtProvider::TensorRT) {
|
||
if (trt_cfg.input_name.empty()) {
|
||
Ort::SessionOptions probe_opts;
|
||
probe_opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL);
|
||
Ort::Session probe(env_, model_path.c_str(), probe_opts);
|
||
Ort::AllocatorWithDefaultOptions alloc;
|
||
trt_cfg.input_name = probe.GetInputNameAllocated(0, alloc).get();
|
||
}
|
||
if (trt_cfg.profile_min.empty()) {
|
||
const std::string shape =
|
||
"1x3x" + std::to_string(kInputH) + "x" + std::to_string(kInputW);
|
||
trt_cfg.profile_min = shape;
|
||
trt_cfg.profile_opt = shape;
|
||
trt_cfg.profile_max = shape;
|
||
}
|
||
}
|
||
|
||
apply_ort_provider(opts, provider, "SCRFDDecoder", trt_cfg);
|
||
|
||
session_ = std::make_unique<Ort::Session>(env_, model_path.c_str(), opts);
|
||
|
||
Ort::AllocatorWithDefaultOptions alloc;
|
||
auto in_name = session_->GetInputNameAllocated(0, alloc);
|
||
input_name_ = in_name.get();
|
||
|
||
const size_t n_out = session_->GetOutputCount();
|
||
if (n_out % 3 != 0 || n_out < 9 || n_out > 12)
|
||
throw std::runtime_error(
|
||
"[SCRFDDecoder] expected 9 or 12 outputs (kps-variant model), got "
|
||
+ std::to_string(n_out));
|
||
fmc_ = static_cast<int>(n_out / 3);
|
||
|
||
for (size_t i = 0; i < n_out; ++i) {
|
||
auto name = session_->GetOutputNameAllocated(i, alloc);
|
||
out_name_storage_.emplace_back(name.get());
|
||
}
|
||
for (auto& s : out_name_storage_)
|
||
out_name_ptrs_.push_back(s.c_str());
|
||
|
||
// Reject non-SCRFD models (e.g. YuNet, which also has 12 outputs and so
|
||
// passes the count check above, but is encoded entirely differently).
|
||
// Cross-check by output channel count: SCRFD's three groups of fmc_
|
||
// outputs encode scores (1ch), bboxes (4ch) and 5-point kps (10ch).
|
||
// YuNet exports loc/conf/iou with 14/2/1 channels, so this trips.
|
||
const int expected_last[3] = {1, 4, 10};
|
||
for (size_t gi = 0; gi < 3; ++gi) {
|
||
for (int si = 0; si < fmc_; ++si) {
|
||
const size_t oi = gi * fmc_ + si;
|
||
auto shape = session_->GetOutputTypeInfo(oi)
|
||
.GetTensorTypeAndShapeInfo().GetShape();
|
||
if (shape.empty() || shape.back() != expected_last[gi]) {
|
||
throw std::runtime_error(
|
||
"[SCRFDDecoder] model does not look like InsightFace SCRFD: "
|
||
"output '" + out_name_storage_[oi] + "' last-dim is "
|
||
+ std::to_string(shape.empty() ? -1 : shape.back())
|
||
+ ", expected " + std::to_string(expected_last[gi])
|
||
+ ". Hint: pass scrfd_500m_bnkps.onnx, not yunet/*.onnx.");
|
||
}
|
||
}
|
||
}
|
||
|
||
std::cerr << "[SCRFDDecoder] loaded: " << model_path << "\n";
|
||
}
|
||
|
||
// Thread-safe: ORT Run() is safe for concurrent calls on the same Session.
|
||
std::vector<DetectedFace> detect(const cv::Mat& img) const {
|
||
// Letterbox to 640×640: uniform scale (preserves aspect ratio) + pad
|
||
// shorter side with constant grey. Stretching to 640×640 (the prior
|
||
// behaviour) distorts faces non-uniformly and degrades landmark
|
||
// localisation — matters most for portrait gallery images and 16:9
|
||
// video frames alike. Coordinates are mapped back via inverse scale +
|
||
// pad-offset below.
|
||
const float scale = std::min(static_cast<float>(kInputW) / img.cols,
|
||
static_cast<float>(kInputH) / img.rows);
|
||
const int new_w = static_cast<int>(std::round(img.cols * scale));
|
||
const int new_h = static_cast<int>(std::round(img.rows * scale));
|
||
const int pad_x = (kInputW - new_w) / 2;
|
||
const int pad_y = (kInputH - new_h) / 2;
|
||
|
||
cv::Mat resized;
|
||
cv::resize(img, resized, {new_w, new_h}, 0, 0, cv::INTER_LINEAR);
|
||
cv::Mat letterboxed(kInputH, kInputW, img.type(),
|
||
cv::Scalar(114, 114, 114));
|
||
resized.copyTo(letterboxed(cv::Rect(pad_x, pad_y, new_w, new_h)));
|
||
|
||
// BGR→RGB swap + normalize to [-1,1] → NCHW float32 blob
|
||
cv::Mat blob = cv::dnn::blobFromImage(
|
||
letterboxed, 1.0 / 128.0, {kInputW, kInputH},
|
||
cv::Scalar(127.5f, 127.5f, 127.5f),
|
||
/*swapRB=*/true, /*crop=*/false, CV_32F);
|
||
|
||
const std::array<int64_t, 4> in_shape = {1, 3, kInputH, kInputW};
|
||
auto mem = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
|
||
auto in_tensor = Ort::Value::CreateTensor<float>(
|
||
mem, blob.ptr<float>(), blob.total(),
|
||
in_shape.data(), in_shape.size());
|
||
|
||
const char* in_name_c = input_name_.c_str();
|
||
auto outs = session_->Run(
|
||
Ort::RunOptions{nullptr},
|
||
&in_name_c, &in_tensor, 1,
|
||
out_name_ptrs_.data(), out_name_ptrs_.size());
|
||
|
||
std::vector<cv::Rect2d> raw_boxes;
|
||
std::vector<float> raw_scores;
|
||
std::vector<std::array<cv::Point2f, 5>> raw_kps;
|
||
|
||
for (int si = 0; si < fmc_; ++si) {
|
||
const int stride = kAllStrides[si];
|
||
const int fh = kInputH / stride;
|
||
const int fw = kInputW / stride;
|
||
|
||
const float* s = outs[si].GetTensorData<float>();
|
||
const float* b = outs[fmc_ + si].GetTensorData<float>();
|
||
const float* k = outs[fmc_ * 2 + si].GetTensorData<float>();
|
||
|
||
for (int r = 0; r < fh; ++r) {
|
||
for (int c = 0; c < fw; ++c) {
|
||
for (int a = 0; a < kAnchors; ++a) {
|
||
const int idx = (r * fw + c) * kAnchors + a;
|
||
const float score = s[idx];
|
||
if (score < conf_threshold_) continue;
|
||
|
||
const float cx = static_cast<float>(c * stride);
|
||
const float cy = static_cast<float>(r * stride);
|
||
|
||
// Decode in letterboxed network space, then un-pad +
|
||
// un-scale to original image coordinates.
|
||
const auto to_img_x = [&](float v) { return (v - pad_x) / scale; };
|
||
const auto to_img_y = [&](float v) { return (v - pad_y) / scale; };
|
||
|
||
const float x1 = to_img_x(cx - b[idx*4+0] * stride);
|
||
const float y1 = to_img_y(cy - b[idx*4+1] * stride);
|
||
const float x2 = to_img_x(cx + b[idx*4+2] * stride);
|
||
const float y2 = to_img_y(cy + b[idx*4+3] * stride);
|
||
raw_boxes.push_back({(double)x1, (double)y1,
|
||
(double)(x2-x1), (double)(y2-y1)});
|
||
raw_scores.push_back(score);
|
||
|
||
std::array<cv::Point2f, 5> lms;
|
||
for (int p = 0; p < 5; ++p)
|
||
lms[p] = {to_img_x(cx + k[idx*10+p*2 ] * stride),
|
||
to_img_y(cy + k[idx*10+p*2+1] * stride)};
|
||
raw_kps.push_back(lms);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
std::vector<int> keep;
|
||
cv::dnn::NMSBoxes(raw_boxes, raw_scores, conf_threshold_, nms_threshold_, keep);
|
||
|
||
const float img_w = static_cast<float>(img.cols);
|
||
const float img_h = static_cast<float>(img.rows);
|
||
|
||
std::vector<DetectedFace> faces;
|
||
faces.reserve(keep.size());
|
||
for (int i : keep) {
|
||
const auto& rb = raw_boxes[i];
|
||
DetectedFace f;
|
||
const float x = std::max(0.f, (float)rb.x);
|
||
const float y = std::max(0.f, (float)rb.y);
|
||
f.bbox = {x, y,
|
||
std::min((float)rb.width, img_w - x),
|
||
std::min((float)rb.height, img_h - y)};
|
||
f.confidence = raw_scores[i];
|
||
f.landmarks = raw_kps[i];
|
||
faces.push_back(f);
|
||
}
|
||
return faces;
|
||
}
|
||
|
||
private:
|
||
float conf_threshold_;
|
||
float nms_threshold_;
|
||
int fmc_{3};
|
||
Ort::Env env_{ORT_LOGGING_LEVEL_WARNING, "scrfd"};
|
||
std::unique_ptr<Ort::Session> session_;
|
||
std::string input_name_;
|
||
std::vector<std::string> out_name_storage_;
|
||
std::vector<const char*> out_name_ptrs_;
|
||
};
|