ANSCORE/modules/ANSFR/ARCFaceRT.h

#ifndef ARCFACERT_H
#define ARCFACERT_H
#pragma once
#include "ANSFRCommon.h"
//#include "fastdeploy/vision.h"
#include "hnswlib/hnswlib.h"
#include <faiss/IndexFlat.h>
#include "engine.h"

namespace ANSCENTER {

    class ArcFace : public ANSFRBase {
    public:
        virtual bool Initialize(std::string licenseKey,
            ModelConfig modelConfig,
            const std::string& modelZipFilePath,
            const std::string& modelZipPassword,
            std::string& labelMap) override;

        virtual bool LoadModel(const std::string& modelZipFilePath,
            const std::string& modelZipPassword) override;

        bool OptimizeModel(bool fp16, std::string& optimizedModelFolder);

        // Single face feature (uses RunArcFace)
        std::vector<float> Feature(const cv::Mat& image, const ANSCENTER::Object& bBox);

        // Main entry: run recognition for all faces in bBox
        std::vector<FaceResultObject> Match(const cv::Mat& input,
            const std::vector<ANSCENTER::Object>& bBox,
            const std::map<std::string, std::string>& userDict);

        cv::Mat GetCropFace(const cv::Mat& input, const ANSCENTER::Object& bBox);

        void Init();

        void AddEmbedding(const std::string& className, float embedding[]);
        void AddEmbedding(const std::string& className, const std::vector<float>& embedding);

        bool UpdateParamater(double knownPersonThreshold) {
            _modelConfig.unknownPersonThreshold = knownPersonThreshold;
            m_knownPersonThresh = _modelConfig.unknownPersonThreshold;
            return true;
        }

        ~ArcFace();
        bool Destroy();

    private:
        bool LoadEngine(const std::string onnxModelPath, bool engineOptimisation = true);

        // Batched forward: one embedding per Object.mask
        std::vector<std::vector<float>> Forward(const cv::Mat& input,
            const std::vector<ANSCENTER::Object>& outputBbox);

        // Search embeddings in FAISS index
        std::tuple<std::vector<std::string>, std::vector<float>>
            SearchForFaces(const std::vector<std::vector<float>>& detectedEmbeddings);

        // Single-face inference (kept for Feature)
        std::vector<float> RunArcFace(const cv::Mat& input);

        // New: batched inference helper
        std::vector<std::vector<float>> RunArcFaceBatch(const std::vector<cv::Mat>& faceROIs);

    protected:
        std::vector<std::string> classNames;
        ModelConfig _modelConfig;
        std::string _modelFilePath;
        float m_knownPersonThresh = 0.35f;

        const int FACE_WIDTH = 112;
        const int FACE_HEIGHT = 112;
        const int FACE_EMBEDDING_SIZE = 512;

        std::unique_ptr<faiss::IndexFlatL2> faiss_index;
        ANSCENTER::Options m_options;
        std::recursive_mutex _mutex;
        Engine<float> m_trtEngine;

        float m_ratio = 1.0f;
        float m_imgWidth = 0.0f;
        float m_imgHeight = 0.0f;

        const std::array<float, 3> SUB_VALS{ 0.5f, 0.5f, 0.5f };
        const std::array<float, 3> DIV_VALS{ 0.5f, 0.5f, 0.5f };
        const bool NORMALIZE = true;
    };
    // Using FastDeploy Enging to perform facial recognition
  //  class ArcFace :public ANSFRBase {
  //  public:
  //      virtual bool Initialize(std::string licenseKey, ModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, std::string& labelMap) override;
  //      virtual bool LoadModel(const std::string& modelZipFilePath, const std::string& modelZipPassword)override;
  //      bool OptimizeModel(bool fp16, std::string& optimizedModelFolder);
  //      std::vector<float>   Feature(const cv::Mat& input, ANSCENTER::Object bBox); // Run inference and get embedding information from a cropped image
  //      std::vector<FaceResultObject>  Match(const cv::Mat& input, std::vector<ANSCENTER::Object> bBox, std::map<std::string, std::string>userDict); // Run inference and get embedding information from a cropped image (the first bbox)
  //      cv::Mat GetCropFace(const cv::Mat& input, ANSCENTER::Object bBox);
  //      void Init();
  //      void AddEmbedding(const std::string& className, float embedding[]);
  //      void AddEmbedding(const std::string& className, const std::vector<float>& embedding);
  //      bool    UpdateParamater(double knownPersonThreshold) {
  //          _modelConfig.unknownPersonThreshold = knownPersonThreshold;
  //          _modelConfig.unknownPersonThreshold = knownPersonThreshold;
  //          m_knownPersonThresh = _modelConfig.unknownPersonThreshold;
  //          return true;
  //      }
		//~ArcFace();
  //      bool Destroy();
  //  private:
  //      bool LoadEngine(const std::string engineFile, bool engineOptimisation = true);
  //      std::vector<std::vector<float>> Forward(const cv::Mat& input, std::vector<ANSCENTER::Object> outputBbox);
  //      std::tuple<std::vector<std::string>, std::vector<float>>  SearchForFaces(const std::vector<std::vector<float>>& detectedEmbeddings);
		//std::vector<float> RunArcFace(const cv::Mat& input);
  //  protected:
  //      std::vector<std::string> classNames;
  //      ModelConfig         _modelConfig;
  //      std::string         _modelFilePath;
  //      float               m_knownPersonThresh;
  //      const int FACE_WIDTH = 112;
  //      const int FACE_HEIGHT = 112;
  //      const int FACE_EMBEDDING_SIZE = 512;
  //      std::unique_ptr<faiss::IndexFlatL2> faiss_index;  // Use shared_ptr
  //      ANSCENTER::Options m_options;
  //      std::recursive_mutex		_mutex;
  //      Engine<float> m_trtEngine;
  //      float m_ratio = 1;
  //      float m_imgWidth = 0;
  //      float m_imgHeight = 0;
  //      const std::array<float, 3> SUB_VALS{ 0.5f, 0.5f, 0.5f };
  //      const std::array<float, 3> DIV_VALS{ 0.5f, 0.5f, 0.5f };
  //      const bool NORMALIZE = true;
  //  };
}
#endif