Initial commit: scene-actor-extraction pipeline
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
This commit is contained in:
@@ -0,0 +1,66 @@
|
||||
#pragma once
|
||||
#include "arcface_embedder.hpp"
|
||||
#include "config.hpp"
|
||||
#include "ort_provider.hpp"
|
||||
#include "trt_arcface_embedder.hpp"
|
||||
|
||||
#include <memory>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
// ── EmbedderFunc ──────────────────────────────────────────────────────────────
|
||||
// KPN node: runs ArcFace on every 112×112 crop in an AlignedSceneFrame,
|
||||
// producing one L2-normalised 512-dim embedding per face.
|
||||
//
|
||||
// Backend selection:
|
||||
// --arcface-engine <path> → TrtArcFaceEmbedder (raw TensorRT, no ORT)
|
||||
// otherwise → ArcFaceEmbedder (ONNX Runtime, picks best EP)
|
||||
//
|
||||
// All crops in one frame are batched into a single forward pass (capped at
|
||||
// embed_batch_size). The backends serialise themselves; we only call them
|
||||
// from the single embedder thread.
|
||||
|
||||
struct EmbedderFunc {
|
||||
static constexpr std::string_view label() { return "embedder"; }
|
||||
|
||||
explicit EmbedderFunc(const Config& cfg, OrtProvider provider)
|
||||
: batch_size_(std::max(1, cfg.embed_batch_size))
|
||||
{
|
||||
if (!cfg.arcface_engine.empty()) {
|
||||
trt_ = std::make_unique<TrtArcFaceEmbedder>(cfg.arcface_engine);
|
||||
if (trt_->max_batch() < static_cast<int>(batch_size_))
|
||||
throw std::runtime_error(
|
||||
"embed_batch_size " + std::to_string(batch_size_) +
|
||||
" exceeds engine max_batch " + std::to_string(trt_->max_batch()) +
|
||||
" — rebuild engine with EMBED_BATCH=" + std::to_string(batch_size_));
|
||||
} else {
|
||||
ort_ = std::make_unique<ArcFaceEmbedder>(
|
||||
cfg.arcface_model, provider, cfg.trt, cfg.embed_batch_size);
|
||||
}
|
||||
}
|
||||
|
||||
EmbeddedSceneFrame operator()(AlignedSceneFrame af) {
|
||||
if (af.source.eof || af.crops.empty())
|
||||
return {std::move(af.source), {}, {}, {}};
|
||||
|
||||
const auto& crops = af.crops;
|
||||
std::vector<Embedding> embeddings;
|
||||
embeddings.reserve(crops.size());
|
||||
for (size_t i = 0; i < crops.size(); i += batch_size_) {
|
||||
const size_t end = std::min(i + batch_size_, crops.size());
|
||||
std::vector<cv::Mat> chunk_crops(crops.begin() + i, crops.begin() + end);
|
||||
auto chunk = trt_ ? trt_->embed(chunk_crops) : ort_->embed(chunk_crops);
|
||||
embeddings.insert(embeddings.end(), chunk.begin(), chunk.end());
|
||||
}
|
||||
|
||||
return {std::move(af.source),
|
||||
std::move(af.faces),
|
||||
std::move(af.crops),
|
||||
std::move(embeddings)};
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<ArcFaceEmbedder> ort_;
|
||||
std::unique_ptr<TrtArcFaceEmbedder> trt_;
|
||||
size_t batch_size_;
|
||||
};
|
||||
Reference in New Issue
Block a user