692 lines
25 KiB
C++
692 lines
25 KiB
C++
#include "ARCFaceRT.h"
|
|
#include "NvOnnxParser.h"
|
|
|
|
namespace ANSCENTER {
|
|
bool ArcFace::Initialize(std::string licenseKey,
|
|
ModelConfig modelConfig,
|
|
const std::string& modelZipFilePath,
|
|
const std::string& modelZipPassword,
|
|
std::string& labelMap) {
|
|
bool result = ANSFRBase::Initialize(licenseKey, modelConfig, modelZipFilePath, modelZipPassword, labelMap);
|
|
if (!result) return false;
|
|
|
|
try {
|
|
_modelConfig = modelConfig;
|
|
_modelConfig.modelType = ModelType::FACERECOGNIZE;
|
|
_modelConfig.detectionType = DetectionType::FACERECOGNIZER;
|
|
|
|
m_knownPersonThresh = _modelConfig.unknownPersonThreshold;
|
|
if (m_knownPersonThresh == 0.0f) m_knownPersonThresh = 0.35f;
|
|
|
|
std::string onnxfile50 = CreateFilePath(_modelFolder, "ansfacerecognizer50.onnx");
|
|
if (std::filesystem::exists(onnxfile50)) {
|
|
_modelFilePath = onnxfile50;
|
|
_logger.LogDebug("ArcFace::Initialize. Loading arcface50 weight", _modelFilePath, __FILE__, __LINE__);
|
|
}
|
|
else {
|
|
std::string onnxfile = CreateFilePath(_modelFolder, "ansfacerecognizer.onnx");
|
|
if (std::filesystem::exists(onnxfile)) {
|
|
_modelFilePath = onnxfile;
|
|
_logger.LogDebug("ArcFace::Initialize. Loading arcface weight", _modelFilePath, __FILE__, __LINE__);
|
|
}
|
|
else {
|
|
_logger.LogError("ArcFace::Initialize. Model arcface.onnx file does not exist", _modelFilePath, __FILE__, __LINE__);
|
|
return false;
|
|
}
|
|
}
|
|
|
|
// Configure engine with batch support
|
|
m_options.precision = ANSCENTER::Precision::FP32;
|
|
m_options.optBatchSize = 8; // expected typical batch
|
|
m_options.maxBatchSize = 32; // maximum number of faces per frame you want
|
|
m_options.calibrationBatchSize = 8;
|
|
m_options.deviceIndex = 0;
|
|
m_trtEngine.UpdateOptions(m_options);
|
|
|
|
if (FileExist(_modelFilePath)) {
|
|
bool succ = m_trtEngine.buildLoadNetwork(_modelFilePath, SUB_VALS, DIV_VALS, NORMALIZE);
|
|
if (!succ) {
|
|
_logger.LogError("ArcFace::Initialize. Unable to build or load TensorRT engine.",
|
|
_modelFilePath, __FILE__, __LINE__);
|
|
return false;
|
|
}
|
|
}
|
|
else {
|
|
_logger.LogError("ArcFace::Initialize. Model file does not exist",
|
|
_modelFilePath, __FILE__, __LINE__);
|
|
return false;
|
|
}
|
|
|
|
Init();
|
|
_isInitialized = true;
|
|
return true;
|
|
}
|
|
catch (const std::exception& e) {
|
|
_logger.LogFatal("ArcFace::Initialize", e.what(), __FILE__, __LINE__);
|
|
return false;
|
|
}
|
|
}
|
|
|
|
bool ArcFace::LoadModel(const std::string& modelZipFilePath, const std::string& modelZipPassword) {
|
|
try {
|
|
bool result = ANSFRBase::LoadModel(modelZipFilePath, modelZipPassword);
|
|
if (!result) return false;
|
|
|
|
std::string onnxfile50 = CreateFilePath(_modelFolder, "ansfacerecognizer50.onnx");
|
|
if (std::filesystem::exists(onnxfile50)) {
|
|
_modelFilePath = onnxfile50;
|
|
_logger.LogDebug("ArcFace::LoadModel. Loading arcface50 weight", _modelFilePath, __FILE__, __LINE__);
|
|
}
|
|
else {
|
|
std::string onnxfile = CreateFilePath(_modelFolder, "ansfacerecognizer.onnx");
|
|
if (std::filesystem::exists(onnxfile)) {
|
|
_modelFilePath = onnxfile;
|
|
_logger.LogDebug("ArcFace::LoadModel. Loading arcface weight", _modelFilePath, __FILE__, __LINE__);
|
|
}
|
|
else {
|
|
_logger.LogError("ArcFace::LoadModel. Model arcface.onnx file does not exist",
|
|
_modelFilePath, __FILE__, __LINE__);
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
catch (const std::exception& e) {
|
|
_logger.LogFatal("ArcFace50::LoadModel", e.what(), __FILE__, __LINE__);
|
|
return false;
|
|
}
|
|
}
|
|
|
|
bool ArcFace::OptimizeModel(bool fp16, std::string& optimizedModelFolder) {
|
|
if (!FileExist(_modelFilePath)) {
|
|
optimizedModelFolder = "";
|
|
return false;
|
|
}
|
|
optimizedModelFolder = GetParentFolder(_modelFilePath);
|
|
|
|
m_options.optBatchSize = 8;
|
|
m_options.maxBatchSize = 32;
|
|
m_options.engineFileDir = optimizedModelFolder;
|
|
m_options.precision = Precision::FP32;
|
|
|
|
Engine<float> engine(m_options);
|
|
|
|
auto succ = engine.buildWithRetry(_modelFilePath, SUB_VALS, DIV_VALS, NORMALIZE);
|
|
if (!succ) {
|
|
const std::string errMsg =
|
|
"Error: Unable to build the TensorRT engine. Try increasing TensorRT log severity to kVERBOSE.";
|
|
_logger.LogError("ArcFace::OptimizeModel", errMsg, __FILE__, __LINE__);
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
std::vector<float> ArcFace::Feature(const cv::Mat& image, const ANSCENTER::Object& bBox) {
|
|
std::vector<float> embedding;
|
|
|
|
// Early validation before locking
|
|
if (image.empty()) {
|
|
return embedding;
|
|
}
|
|
|
|
if (image.cols < 10 || image.rows < 10) {
|
|
return embedding;
|
|
}
|
|
|
|
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
|
|
|
try {
|
|
return RunArcFace(bBox.mask);
|
|
}
|
|
catch (const std::exception& e) {
|
|
_logger.LogFatal("ArcFace::Feature", e.what(), __FILE__, __LINE__);
|
|
return std::vector<float>();
|
|
}
|
|
}
|
|
std::vector<FaceResultObject> ArcFace::Match(
|
|
const cv::Mat& input,
|
|
const std::vector<ANSCENTER::Object>& bBox,
|
|
const std::map<std::string, std::string>& userDict) {
|
|
|
|
std::vector<FaceResultObject> resultObjects;
|
|
|
|
// Early validation before locking
|
|
if (input.empty()) {
|
|
return resultObjects;
|
|
}
|
|
|
|
if (input.cols < 10 || input.rows < 10) {
|
|
return resultObjects;
|
|
}
|
|
|
|
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
|
|
|
if (!_isInitialized) {
|
|
_logger.LogError("ArcFace::Match", "Model is not initialized", __FILE__, __LINE__);
|
|
return resultObjects;
|
|
}
|
|
|
|
try {
|
|
// Get embeddings
|
|
std::vector<std::vector<float>> detectedEmbeddings = Forward(input, bBox);
|
|
|
|
// Search for matches
|
|
std::vector<std::string> names;
|
|
std::vector<float> sims;
|
|
std::tie(names, sims) = SearchForFaces(detectedEmbeddings);
|
|
|
|
if (names.empty()) {
|
|
_logger.LogError("ArcFace::Match", "No face is match", __FILE__, __LINE__);
|
|
return resultObjects;
|
|
}
|
|
|
|
// Pre-reserve result space
|
|
const size_t resultCount = std::min(names.size(), bBox.size());
|
|
resultObjects.reserve(resultCount);
|
|
|
|
// Build result objects
|
|
for (size_t i = 0; i < resultCount; ++i) {
|
|
FaceResultObject resultObject;
|
|
|
|
// Determine if face is known or unknown
|
|
const bool isUnknown = (sims[i] > m_knownPersonThresh);
|
|
|
|
if (isUnknown) {
|
|
resultObject.isUnknown = true;
|
|
resultObject.userId = "0";
|
|
resultObject.userName = "Unknown";
|
|
resultObject.confidence = 1.0f;
|
|
}
|
|
else {
|
|
resultObject.isUnknown = false;
|
|
resultObject.userId = names[i];
|
|
|
|
// Safe map lookup with fallback
|
|
auto it = userDict.find(names[i]);
|
|
resultObject.userName = (it != userDict.end()) ? it->second : names[i];
|
|
|
|
resultObject.confidence = 1.0f - sims[i];
|
|
}
|
|
|
|
resultObject.similarity = sims[i];
|
|
|
|
// Copy bounding box and additional data
|
|
resultObject.box = bBox[i].box;
|
|
resultObject.mask = bBox[i].mask;
|
|
resultObject.cameraId = bBox[i].cameraId;
|
|
resultObject.trackId = bBox[i].trackId;
|
|
resultObject.polygon = bBox[i].polygon;
|
|
resultObject.kps = bBox[i].kps;
|
|
resultObject.extraInformation = bBox[i].extraInfo;
|
|
|
|
resultObjects.push_back(std::move(resultObject));
|
|
}
|
|
|
|
return resultObjects;
|
|
}
|
|
catch (const std::exception& e) {
|
|
_logger.LogFatal("ArcFace::Match", e.what(), __FILE__, __LINE__);
|
|
return resultObjects;
|
|
}
|
|
}
|
|
|
|
cv::Mat ArcFace::GetCropFace(const cv::Mat& input, const ANSCENTER::Object& bBox) {
|
|
try {
|
|
std::vector<ANSCENTER::Object> outputBbox;
|
|
outputBbox.reserve(1);
|
|
outputBbox.push_back(bBox);
|
|
|
|
std::vector<CroppedFace> crFaces;
|
|
ANSFRHelper::GetCroppedFaces(input, outputBbox, 112, 112, crFaces);
|
|
|
|
if (crFaces.empty()) {
|
|
return cv::Mat();
|
|
}
|
|
|
|
return crFaces[0].faceMat;
|
|
}
|
|
catch (const std::exception& e) {
|
|
_logger.LogFatal("ArcFace::GetCropFace", e.what(), __FILE__, __LINE__);
|
|
return cv::Mat();
|
|
}
|
|
}
|
|
|
|
bool ArcFace::LoadEngine(const std::string onnxModelPath, bool engineOptimisation) {
|
|
try {
|
|
if (!FileExist(onnxModelPath)) {
|
|
_logger.LogError("ArcFace::LoadEngine", "Cannot find the raw ONNX model file.", __FILE__, __LINE__);
|
|
return false;
|
|
}
|
|
|
|
m_options.precision = ANSCENTER::Precision::FP32;
|
|
m_options.optBatchSize = 8;
|
|
m_options.maxBatchSize = 32;
|
|
m_options.calibrationBatchSize = 8;
|
|
m_options.deviceIndex = 0;
|
|
m_trtEngine.UpdateOptions(m_options);
|
|
|
|
if (FileExist(onnxModelPath)) {
|
|
bool succ = m_trtEngine.buildLoadNetwork(onnxModelPath, SUB_VALS, DIV_VALS, NORMALIZE);
|
|
if (!succ) {
|
|
_logger.LogError("ArcFace::LoadEngine. Unable to build or load TensorRT engine.",
|
|
onnxModelPath, __FILE__, __LINE__);
|
|
return false;
|
|
}
|
|
}
|
|
else {
|
|
_logger.LogError("ArcFace::LoadEngine. Model file does not exist",
|
|
onnxModelPath, __FILE__, __LINE__);
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
catch (const std::exception& e) {
|
|
_logger.LogFatal("ArcFace::LoadEngine", e.what(), __FILE__, __LINE__);
|
|
return false;
|
|
}
|
|
}
|
|
|
|
std::vector<float> ArcFace::RunArcFace(const cv::Mat& inputImage) {
|
|
std::vector<float> embedding;
|
|
|
|
// Early validation before locking
|
|
if (inputImage.empty()) {
|
|
_logger.LogError("ArcFace::RunArcFace", "Input image is empty", __FILE__, __LINE__);
|
|
return embedding;
|
|
}
|
|
|
|
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
|
|
|
try {
|
|
if (!_isInitialized) {
|
|
_logger.LogError("ArcFace::RunArcFace", "Model is not initialized", __FILE__, __LINE__);
|
|
return embedding;
|
|
}
|
|
|
|
// GPU preprocessing pipeline
|
|
cv::cuda::Stream stream;
|
|
cv::cuda::GpuMat d_img;
|
|
|
|
// Upload to GPU
|
|
d_img.upload(inputImage, stream);
|
|
|
|
// Handle grayscale conversion on GPU
|
|
if (inputImage.channels() == 1) {
|
|
cv::cuda::GpuMat d_bgr;
|
|
cv::cuda::cvtColor(d_img, d_bgr, cv::COLOR_GRAY2BGR, 0, stream);
|
|
d_img = d_bgr;
|
|
}
|
|
|
|
// Resize on GPU if needed
|
|
if (inputImage.cols != FACE_WIDTH || inputImage.rows != FACE_HEIGHT) {
|
|
cv::cuda::GpuMat d_resized;
|
|
cv::cuda::resize(d_img, d_resized, cv::Size(FACE_WIDTH, FACE_HEIGHT),
|
|
0, 0, cv::INTER_LINEAR, stream);
|
|
d_img = d_resized;
|
|
}
|
|
|
|
// BGR to RGB conversion on GPU
|
|
cv::cuda::GpuMat d_rgb;
|
|
cv::cuda::cvtColor(d_img, d_rgb, cv::COLOR_BGR2RGB, 0, stream);
|
|
|
|
// Prepare inference inputs
|
|
std::vector<cv::cuda::GpuMat> inputVec;
|
|
inputVec.emplace_back(std::move(d_rgb));
|
|
|
|
std::vector<std::vector<cv::cuda::GpuMat>> inputs;
|
|
inputs.emplace_back(std::move(inputVec));
|
|
|
|
// Run inference
|
|
std::vector<std::vector<std::vector<float>>> featureVectors;
|
|
bool succ = m_trtEngine.runInference(inputs, featureVectors);
|
|
|
|
stream.waitForCompletion();
|
|
|
|
if (!succ) {
|
|
_logger.LogError("ArcFace::RunArcFace", "Failed to run inference.", __FILE__, __LINE__);
|
|
return embedding;
|
|
}
|
|
|
|
if (!featureVectors.empty() && !featureVectors[0].empty()) {
|
|
embedding = std::move(featureVectors[0][0]);
|
|
}
|
|
|
|
return embedding;
|
|
}
|
|
catch (const std::exception& e) {
|
|
_logger.LogFatal("ArcFace::RunArcFace", e.what(), __FILE__, __LINE__);
|
|
return embedding;
|
|
}
|
|
}
|
|
std::vector<std::vector<float>> ArcFace::RunArcFaceBatch(const std::vector<cv::Mat>& faceROIs) {
|
|
std::vector<std::vector<float>> embeddings;
|
|
|
|
try {
|
|
if (!_isInitialized) {
|
|
_logger.LogError("ArcFace::RunArcFaceBatch", "Model is not initialized", __FILE__, __LINE__);
|
|
return embeddings;
|
|
}
|
|
|
|
if (faceROIs.empty()) {
|
|
return embeddings;
|
|
}
|
|
|
|
if (faceROIs.size() > static_cast<size_t>(m_options.maxBatchSize)) {
|
|
_logger.LogError("ArcFace::RunArcFaceBatch",
|
|
"Batch size exceeds maxBatchSize", __FILE__, __LINE__);
|
|
return embeddings;
|
|
}
|
|
|
|
const auto& inputDims = m_trtEngine.getInputDims();
|
|
if (inputDims.empty() || inputDims[0].nbDims < 3) {
|
|
_logger.LogError("ArcFace::RunArcFaceBatch",
|
|
"Invalid engine input dims", __FILE__, __LINE__);
|
|
return embeddings;
|
|
}
|
|
|
|
// Pre-reserve embeddings space
|
|
embeddings.reserve(faceROIs.size());
|
|
|
|
// GPU preprocessing pipeline
|
|
cv::cuda::Stream stream;
|
|
std::vector<cv::cuda::GpuMat> batchGpu;
|
|
batchGpu.reserve(faceROIs.size());
|
|
|
|
const cv::Size targetSize(FACE_WIDTH, FACE_HEIGHT);
|
|
|
|
for (size_t i = 0; i < faceROIs.size(); ++i) {
|
|
const cv::Mat& roi = faceROIs[i];
|
|
|
|
if (roi.empty()) {
|
|
_logger.LogWarn("ArcFace::RunArcFaceBatch",
|
|
"Empty ROI at index " + std::to_string(i) + ", skipping",
|
|
__FILE__, __LINE__);
|
|
continue;
|
|
}
|
|
|
|
// Upload to GPU
|
|
cv::cuda::GpuMat d_img;
|
|
d_img.upload(roi, stream);
|
|
|
|
// Handle grayscale conversion on GPU
|
|
if (roi.channels() == 1) {
|
|
cv::cuda::GpuMat d_bgr;
|
|
cv::cuda::cvtColor(d_img, d_bgr, cv::COLOR_GRAY2BGR, 0, stream);
|
|
d_img = d_bgr;
|
|
}
|
|
|
|
// Resize on GPU if needed
|
|
if (roi.cols != FACE_WIDTH || roi.rows != FACE_HEIGHT) {
|
|
cv::cuda::GpuMat d_resized;
|
|
cv::cuda::resize(d_img, d_resized, targetSize, 0, 0, cv::INTER_LINEAR, stream);
|
|
d_img = d_resized;
|
|
}
|
|
|
|
// BGR to RGB conversion on GPU
|
|
cv::cuda::GpuMat d_rgb;
|
|
cv::cuda::cvtColor(d_img, d_rgb, cv::COLOR_BGR2RGB, 0, stream);
|
|
|
|
batchGpu.emplace_back(std::move(d_rgb));
|
|
}
|
|
|
|
if (batchGpu.empty()) {
|
|
return embeddings;
|
|
}
|
|
|
|
// Prepare inference inputs
|
|
std::vector<std::vector<cv::cuda::GpuMat>> inputs;
|
|
inputs.reserve(1);
|
|
inputs.emplace_back(std::move(batchGpu));
|
|
|
|
// Run inference
|
|
std::vector<std::vector<std::vector<float>>> featureVectors;
|
|
bool succ = m_trtEngine.runInference(inputs, featureVectors);
|
|
|
|
stream.waitForCompletion();
|
|
|
|
if (!succ) {
|
|
_logger.LogError("ArcFace::RunArcFaceBatch", "runInference failed", __FILE__, __LINE__);
|
|
return embeddings;
|
|
}
|
|
|
|
if (featureVectors.empty() || featureVectors[0].empty()) {
|
|
_logger.LogError("ArcFace::RunArcFaceBatch", "Empty featureVectors", __FILE__, __LINE__);
|
|
return embeddings;
|
|
}
|
|
|
|
embeddings = std::move(featureVectors[0]);
|
|
|
|
return embeddings;
|
|
}
|
|
catch (const std::exception& e) {
|
|
_logger.LogFatal("ArcFace::RunArcFaceBatch", e.what(), __FILE__, __LINE__);
|
|
return embeddings;
|
|
}
|
|
}
|
|
|
|
std::vector<std::vector<float>> ArcFace::Forward(const cv::Mat& input,const std::vector<ANSCENTER::Object>& outputBbox)
|
|
{
|
|
std::vector<std::vector<float>> detectedEmbeddings;
|
|
|
|
// Early validation before locking
|
|
if (input.empty()) {
|
|
_logger.LogError("ArcFace::Forward",
|
|
"Input image is empty", __FILE__, __LINE__);
|
|
return detectedEmbeddings;
|
|
}
|
|
|
|
if (outputBbox.empty()) {
|
|
return detectedEmbeddings;
|
|
}
|
|
|
|
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
|
|
|
try {
|
|
// Pre-reserve output space
|
|
detectedEmbeddings.reserve(outputBbox.size());
|
|
|
|
// Collect valid face ROIs
|
|
std::vector<cv::Mat> faceROIs;
|
|
faceROIs.reserve(outputBbox.size());
|
|
|
|
for (const auto& obj : outputBbox) {
|
|
if (!obj.mask.empty()) {
|
|
faceROIs.push_back(obj.mask);
|
|
}
|
|
}
|
|
|
|
if (faceROIs.empty()) {
|
|
return detectedEmbeddings;
|
|
}
|
|
|
|
// Run batch inference
|
|
detectedEmbeddings = RunArcFaceBatch(faceROIs);
|
|
|
|
return detectedEmbeddings;
|
|
}
|
|
catch (const std::exception& e) {
|
|
_logger.LogFatal("ArcFace::Forward", e.what(), __FILE__, __LINE__);
|
|
return detectedEmbeddings;
|
|
}
|
|
}
|
|
|
|
std::tuple<std::vector<std::string>, std::vector<float>>
|
|
ArcFace::SearchForFaces(const std::vector<std::vector<float>>& detectedEmbeddings) {
|
|
std::vector<std::string> detectedUsers;
|
|
std::vector<float> simValues;
|
|
|
|
// Early exit before locking
|
|
if (detectedEmbeddings.empty()) {
|
|
return std::make_tuple(detectedUsers, simValues);
|
|
}
|
|
|
|
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
|
|
|
try {
|
|
// Pre-reserve output space
|
|
detectedUsers.reserve(detectedEmbeddings.size());
|
|
simValues.reserve(detectedEmbeddings.size());
|
|
|
|
if (!classNames.empty() && faiss_index && faiss_index->ntotal > 0) {
|
|
// Determine k based on database size
|
|
const int k = std::min(3, static_cast<int>(faiss_index->ntotal));
|
|
|
|
// Pre-allocate search buffers (reuse across iterations)
|
|
std::vector<faiss::idx_t> indices(k);
|
|
std::vector<float> distances(k);
|
|
std::vector<float> matchEmbedding(faiss_index->d);
|
|
|
|
for (const auto& embedding : detectedEmbeddings) {
|
|
if (embedding.size() != FACE_EMBEDDING_SIZE) {
|
|
detectedUsers.push_back("0");
|
|
simValues.push_back(1.0f);
|
|
continue;
|
|
}
|
|
|
|
// Search FAISS index
|
|
faiss_index->search(1, embedding.data(), k,
|
|
distances.data(), indices.data());
|
|
|
|
// Find best match (minimum distance for L2)
|
|
auto min_it = std::min_element(distances.begin(), distances.end());
|
|
int best_index = static_cast<int>(std::distance(distances.begin(), min_it));
|
|
|
|
// Validate index
|
|
faiss::idx_t id = indices[best_index];
|
|
if (id < 0 || id >= static_cast<faiss::idx_t>(classNames.size())) {
|
|
detectedUsers.push_back("0");
|
|
simValues.push_back(1.0f);
|
|
continue;
|
|
}
|
|
|
|
// Reconstruct embedding and compute similarity
|
|
faiss_index->reconstruct(id, matchEmbedding.data());
|
|
float cosine = CosineSimilarity(embedding, matchEmbedding, false);
|
|
float similarity = 1.0f - cosine;
|
|
|
|
detectedUsers.push_back(classNames[id]);
|
|
simValues.push_back(std::abs(similarity));
|
|
}
|
|
}
|
|
else {
|
|
detectedUsers.assign(detectedEmbeddings.size(), "0");
|
|
simValues.assign(detectedEmbeddings.size(), 1.0f);
|
|
}
|
|
|
|
return std::make_tuple(std::move(detectedUsers), std::move(simValues));
|
|
}
|
|
catch (const std::exception& e) {
|
|
_logger.LogFatal("ArcFace::SearchForFaces", e.what(), __FILE__, __LINE__);
|
|
|
|
// Return appropriate sized vectors on error
|
|
detectedUsers.assign(detectedEmbeddings.size(), "0");
|
|
simValues.assign(detectedEmbeddings.size(), 1.0f);
|
|
return std::make_tuple(detectedUsers, simValues);
|
|
}
|
|
}
|
|
void ArcFace::Init() {
|
|
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
|
|
|
try {
|
|
classNames.clear();
|
|
|
|
if (faiss_index) {
|
|
faiss_index->reset();
|
|
}
|
|
else {
|
|
faiss_index = std::make_unique<faiss::IndexFlatL2>(FACE_EMBEDDING_SIZE);
|
|
}
|
|
|
|
}
|
|
catch (const std::exception& e) {
|
|
_isInitialized = false;
|
|
_logger.LogFatal("ArcFace::Init", e.what(), __FILE__, __LINE__);
|
|
}
|
|
}
|
|
ArcFace::~ArcFace() {
|
|
try {
|
|
Destroy();
|
|
}
|
|
catch (const std::exception& e) {
|
|
// Log but don't throw - exceptions in destructors are dangerous
|
|
_logger.LogError("ArcFace::~ArcFace", e.what(), __FILE__, __LINE__);
|
|
}
|
|
catch (...) {
|
|
// Catch all exceptions to prevent std::terminate
|
|
_logger.LogError("ArcFace::~ArcFace", "Unknown exception during destruction", __FILE__, __LINE__);
|
|
}
|
|
}
|
|
bool ArcFace::Destroy() {
|
|
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
|
|
|
try {
|
|
classNames.clear();
|
|
|
|
if (faiss_index) {
|
|
faiss_index->reset();
|
|
faiss_index.reset();
|
|
}
|
|
|
|
_isInitialized = false;
|
|
|
|
m_trtEngine.clearGpuBuffers();
|
|
return true;
|
|
}
|
|
catch (const std::exception& e) {
|
|
_logger.LogFatal("ArcFace::Destroy", e.what(), __FILE__, __LINE__);
|
|
return false;
|
|
}
|
|
}
|
|
void ArcFace::AddEmbedding(const std::string& className, float embedding[]) {
|
|
// Validate input before locking
|
|
if (!embedding) {
|
|
_logger.LogError("ArcFace::AddEmbedding",
|
|
"Null embedding pointer.", __FILE__, __LINE__);
|
|
return;
|
|
}
|
|
|
|
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
|
|
|
try {
|
|
if (!faiss_index) {
|
|
_logger.LogError("ArcFace::AddEmbedding",
|
|
"FAISS index is not initialized.", __FILE__, __LINE__);
|
|
return;
|
|
}
|
|
|
|
// Direct add without intermediate vector copy
|
|
classNames.push_back(className);
|
|
faiss_index->add(1, embedding);
|
|
|
|
}
|
|
catch (const std::exception& e) {
|
|
_logger.LogFatal("ArcFace::AddEmbedding", e.what(), __FILE__, __LINE__);
|
|
}
|
|
}
|
|
void ArcFace::AddEmbedding(const std::string& className, const std::vector<float>& embedding) {
|
|
// Early validation before locking
|
|
if (embedding.size() != FACE_EMBEDDING_SIZE) {
|
|
_logger.LogError("ArcFace::AddEmbedding",
|
|
"Embedding size does not match expected output dimension of 512.", __FILE__, __LINE__);
|
|
return;
|
|
}
|
|
|
|
std::lock_guard<std::recursive_mutex> lock(_mutex);
|
|
|
|
try {
|
|
if (!faiss_index) {
|
|
_logger.LogError("ArcFace::AddEmbedding",
|
|
"FAISS index is not initialized.", __FILE__, __LINE__);
|
|
return;
|
|
}
|
|
|
|
classNames.push_back(className);
|
|
faiss_index->add(1, embedding.data());
|
|
|
|
}
|
|
catch (const std::exception& e) {
|
|
_logger.LogFatal("ArcFace::AddEmbedding", e.what(), __FILE__, __LINE__);
|
|
}
|
|
}
|
|
|
|
} |