dtourolle/scene-actor-extraction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add AMD support via ort alternative to trt

Browse Source
This commit is contained in:
dtourolle
2026-06-28 11:50:05 +02:00
parent a3ba53ddf7
commit 0ee131a692
27 changed files with 1357 additions and 977 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/nodes/embedder_node.hpp
+17 -25
View File
@@ -1,9 +1,8 @@
#pragma once
#include "arcface_embedder.hpp"
#include "config.hpp"
#include "ort_provider.hpp"
#include "trt_arcface_embedder.hpp"
#include "inference/face_embedder.hpp"
#include <algorithm>
#include <memory>
#include <stdexcept>
#include <string>
@@ -12,31 +11,25 @@
// 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)
// The inference backend (ONNX Runtime or raw TensorRT) is selected at compile
// time; this node talks only to IFaceEmbedder via make_face_embedder(cfg).
//
// 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.
// embed_batch_size). The backend serialises itself; we only call it 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))
explicit EmbedderFunc(const Config& cfg)
: embedder_(make_face_embedder(cfg))
, 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);
}
if (embedder_->max_batch() < static_cast<int>(batch_size_))
throw std::runtime_error(
"embed_batch_size " + std::to_string(batch_size_) +
" exceeds backend max_batch " + std::to_string(embedder_->max_batch()) +
" — rebuild the engine with EMBED_BATCH=" + std::to_string(batch_size_));
}
EmbeddedSceneFrame operator()(AlignedSceneFrame af) {
@@ -49,7 +42,7 @@ struct EmbedderFunc {
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);
auto chunk = embedder_->embed(chunk_crops);
embeddings.insert(embeddings.end(), chunk.begin(), chunk.end());
}
@@ -60,7 +53,6 @@ struct EmbedderFunc {
}
private:
std::unique_ptr<ArcFaceEmbedder> ort_;
std::unique_ptr<TrtArcFaceEmbedder> trt_;
size_t batch_size_;
std::unique_ptr<IFaceEmbedder> embedder_;
size_t batch_size_;
};
src/nodes/face_detector_node.hpp
+12 -22
View File
@@ -1,39 +1,30 @@
#pragma once
#include "scrfd_decoder.hpp"
#include "trt_scrfd_decoder.hpp"
#include "config.hpp"
#include "ort_provider.hpp"
#include "inference/face_detector.hpp"
#include <algorithm>
#include <memory>
#include <string>
// ── FaceDetectorFunc ──────────────────────────────────────────────────────────
// KPN node: runs SCRFD-500MF to detect ALL faces in a frame.
//
// Backend selection:
// --detector-engine <path> → TrtScrfdDecoder (raw TensorRT, no ORT)
// otherwise → SCRFDDecoder (ONNX Runtime)
// The inference backend (ONNX Runtime or raw TensorRT) is selected at compile
// time; this node talks only to IFaceDetector via make_face_detector(cfg).
struct FaceDetectorFunc {
static constexpr std::string_view label() { return "face_detector"; }
explicit FaceDetectorFunc(const Config& cfg, OrtProvider provider)
: max_faces_(cfg.max_faces)
explicit FaceDetectorFunc(const Config& cfg)
: detector_(make_face_detector(cfg))
, max_faces_(cfg.max_faces)
, min_face_px_(cfg.min_face_px)
{
if (!cfg.detector_engine.empty()) {
trt_ = std::make_unique<TrtScrfdDecoder>(
cfg.detector_engine, cfg.detector_conf, cfg.detector_nms);
} else {
ort_ = std::make_unique<SCRFDDecoder>(
cfg.detector_model, cfg.detector_conf, cfg.detector_nms, provider, cfg.trt);
}
}
{}
SceneFrame operator()(Frame f) {
if (f.eof) return {std::move(f), {}};
auto faces = trt_ ? trt_->detect(f.image) : ort_->detect(f.image);
auto faces = detector_->detect(f.image);
// Drop faces below minimum pixel size (too small for reliable ArcFace alignment)
faces.erase(
@@ -54,8 +45,7 @@ struct FaceDetectorFunc {
}
private:
std::unique_ptr<SCRFDDecoder> ort_;
std::unique_ptr<TrtScrfdDecoder> trt_;
int max_faces_{10};
float min_face_px_{40.f};
std::unique_ptr<IFaceDetector> detector_;
int max_faces_{10};
float min_face_px_{40.f};
};
src/nodes/frame_source_node.hpp
+3 -2
View File
@@ -14,7 +14,8 @@
// ── FrameSourceFunc ───────────────────────────────────────────────────────────
// KPN source node: reads a movie file and emits one Frame per sample interval.
//
// Decode backend: FFmpeg with NVDEC (_cuvid) when available, CPU otherwise.
// Decode backend: FFmpeg hwaccel (CUDA/VAAPI, runtime-detected) when
// available, CPU otherwise.
//
// Sampling strategy: seek to the next target timestamp rather than decoding
// every frame, which is fast even for 1-FPS sampling of a 2-hour film.
@@ -39,7 +40,7 @@ struct FrameSourceFunc {
- cfg.start_sec;
int n_frames = static_cast<int>(span_s * cfg.sample_fps);
std::cerr << "[frame_source] decoder=" << decoder_->codec_name()
<< " (" << (decoder_->hw_active() ? "NVDEC" : "CPU") << ")"
<< " (" << decoder_->hw_backend() << ")"
<< " video_fps=" << decoder_->fps()
<< " start=" << cfg.start_sec << "s"
<< (end_sec_ > 0 ? " end=" + std::to_string(end_sec_) + "s" : "")
src/nodes/identity_matcher_node.hpp
+10 -88
View File
@@ -1,16 +1,15 @@
#pragma once
#include "types.hpp"
#include "config.hpp"
#include "inference/similarity.hpp"
#include "gallery/gallery_store.hpp"
#include "gallery/gallery_calibration.hpp"
#include <cublas_v2.h>
#include <cuda_runtime_api.h>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <limits>
#include <memory>
#include <stdexcept>
#include <vector>
@@ -35,25 +34,9 @@
// Gallery scan: the full reference set (tens of thousands of 512-dim
// embeddings) is uploaded to the GPU once at construction time and stays
// resident there. Per frame, only the small query matrix (n_faces x 512) is
// uploaded and a single cublasSgemm computes the full similarity matrix
// (n_faces x N_gallery) in well under a millisecond — far faster than any
// CPU GEMM or scalar scan, making gallery-side pruning unnecessary.
namespace identity_matcher_detail {
struct CudaError : std::runtime_error {
using std::runtime_error::runtime_error;
};
inline void check_cuda(cudaError_t e, const char* what) {
if (e != cudaSuccess)
throw CudaError(std::string(what) + ": " + cudaGetErrorString(e));
}
inline void check_cublas(cublasStatus_t s, const char* what) {
if (s != CUBLAS_STATUS_SUCCESS)
throw CudaError(std::string(what) + ": cublas error " + std::to_string(s));
}
} // namespace identity_matcher_detail
// uploaded and a single SGEMM computes the full similarity matrix in well under
// a millisecond. The GPU math backend (cuBLAS or rocBLAS) lives behind
// ISimilarityEngine (backends/gemm_backend.cpp) and is selected at compile time.
struct IdentityMatcherFunc {
static constexpr std::string_view label() { return "identity_matcher"; }
@@ -69,8 +52,6 @@ struct IdentityMatcherFunc {
, ratio_(cfg.match_ratio)
, ratio_ceil_(cfg.match_ratio_ceil)
{
using namespace identity_matcher_detail;
std::cerr << "[identity_matcher] flattening gallery embeddings...\n";
for (int ai = 0; ai < static_cast<int>(gallery_.actors.size()); ++ai) {
for (const auto& emb : gallery_.actors[ai].embeddings) {
@@ -100,47 +81,15 @@ struct IdentityMatcherFunc {
<< gallery_.actors.size() << " actors, "
<< flat_emb_.size() << " reference embeddings\n";
// Flatten gallery into a contiguous (N x 512) row-major host buffer,
// then upload once. Row-major NxD == column-major DxN, which is the
// layout cublasSgemm wants for the transposed operand below.
std::vector<float> host_gallery(static_cast<size_t>(n_gallery_) * 512);
for (int i = 0; i < n_gallery_; ++i)
std::memcpy(host_gallery.data() + static_cast<size_t>(i) * 512,
flat_emb_[i].data(), 512 * sizeof(float));
check_cuda(cudaMalloc(reinterpret_cast<void**>(&d_gallery_), host_gallery.size() * sizeof(float)),
"cudaMalloc gallery");
check_cuda(cudaMemcpy(d_gallery_, host_gallery.data(),
host_gallery.size() * sizeof(float),
cudaMemcpyHostToDevice),
"cudaMemcpy gallery H2D");
check_cuda(cudaMalloc(reinterpret_cast<void**>(&d_query_), static_cast<size_t>(kMaxFaces) * 512 * sizeof(float)),
"cudaMalloc query");
check_cuda(cudaMalloc(reinterpret_cast<void**>(&d_sims_), static_cast<size_t>(kMaxFaces) * n_gallery_ * sizeof(float)),
"cudaMalloc sims");
check_cuda(cudaStreamCreate(&stream_), "cudaStreamCreate");
check_cublas(cublasCreate(&handle_), "cublasCreate");
check_cublas(cublasSetStream(handle_, stream_), "cublasSetStream");
host_sims_.resize(static_cast<size_t>(kMaxFaces) * n_gallery_);
std::cerr << "[identity_matcher] gallery resident on GPU ("
<< (host_gallery.size() * sizeof(float)) / (1024 * 1024) << " MiB)\n";
}
~IdentityMatcherFunc() {
if (d_gallery_) cudaFree(d_gallery_);
if (d_query_) cudaFree(d_query_);
if (d_sims_) cudaFree(d_sims_);
if (handle_) cublasDestroy(handle_);
if (stream_) cudaStreamDestroy(stream_);
sim_engine_ = make_similarity_engine(host_gallery.data(), n_gallery_, kMaxFaces);
}
MatchedSceneFrame operator()(TrackedSceneFrame tf) {
using namespace identity_matcher_detail;
if (tf.source.eof) return {std::move(tf.source), {}};
const int n_faces = static_cast<int>(tf.embeddings.size());
@@ -151,37 +100,18 @@ struct IdentityMatcherFunc {
if (n_faces > kMaxFaces)
throw std::runtime_error("identity_matcher: n_faces exceeds kMaxFaces");
// Build the (n_faces x 512) query matrix (row-major == col-major 512 x n_faces).
std::vector<float> host_query(static_cast<size_t>(n_faces) * 512);
for (int fi = 0; fi < n_faces; ++fi) {
std::memcpy(host_query.data() + static_cast<size_t>(fi) * 512,
tf.embeddings[fi].data(), 512 * sizeof(float));
}
check_cuda(cudaMemcpyAsync(d_query_, host_query.data(),
host_query.size() * sizeof(float),
cudaMemcpyHostToDevice, stream_),
"cudaMemcpy query H2D");
// S (N_gallery x n_faces) col-major = G(512 x N_gallery)^T * Q(512 x n_faces)
// i.e. S[g + f*N_gallery] = cosine_similarity(gallery[g], query[f]).
const float alpha = 1.f, beta = 0.f;
check_cublas(cublasSgemm(handle_, CUBLAS_OP_T, CUBLAS_OP_N,
n_gallery_, n_faces, 512,
&alpha, d_gallery_, 512, d_query_, 512,
&beta, d_sims_, n_gallery_),
"cublasSgemm");
check_cuda(cudaMemcpyAsync(host_sims_.data(), d_sims_,
static_cast<size_t>(n_gallery_) * n_faces * sizeof(float),
cudaMemcpyDeviceToHost, stream_),
"cudaMemcpy sims D2H");
check_cuda(cudaStreamSynchronize(stream_), "cudaStreamSynchronize");
// S (N_gallery × n_faces) col-major: face fi's gallery sims at sims + fi*n_gallery.
const float* host_sims = sim_engine_->compute(host_query.data(), n_faces);
for (int fi = 0; fi < n_faces; ++fi) {
const float* sims = host_sims_.data() + static_cast<size_t>(fi) * n_gallery_;
const float* sims = host_sims + static_cast<size_t>(fi) * n_gallery_;
// Per-actor best cosine similarity (max over that actor's reference embeddings)
std::vector<float> best_sim(gallery_.actors.size(),
-std::numeric_limits<float>::max());
for (int ei = 0; ei < n_gallery_; ++ei) {
@@ -190,7 +120,6 @@ struct IdentityMatcherFunc {
if (sim > best_sim[ai]) best_sim[ai] = sim;
}
// Find best and second-best actor by similarity
int best_actor = -1;
int second_actor = -1;
float best_s = -std::numeric_limits<float>::max();
@@ -240,7 +169,6 @@ struct IdentityMatcherFunc {
? cal_.probability(best_s, log_prior_odds_)
: best_s;
}
// actor_idx == -1, name == "" → unknown face
actors.push_back(std::move(ia));
}
@@ -260,11 +188,5 @@ private:
std::vector<int> flat_actor_;
int n_gallery_{0};
float* d_gallery_{nullptr}; // (n_gallery_ x 512), resident for the lifetime of this node
float* d_query_{nullptr}; // (kMaxFaces x 512)
float* d_sims_{nullptr}; // (n_gallery_ x kMaxFaces), column-major
std::vector<float> host_sims_;
cudaStream_t stream_{nullptr};
cublasHandle_t handle_{nullptr};
std::unique_ptr<ISimilarityEngine> sim_engine_;
};