#ifndef _FACE_MASK_H
#define _FACE_MASK_H

#include <vector>
#include <string>
#include "ai_utils.h"
#include "face_detection.h"

using namespace std;

typedef struct FaceMaskInfo
{
    std::string label; 
    float score;
} FaceMaskInfo;

class FaceMask
{
public:
    FaceMask(std::string &kmodel_file,float mask_thresh, FrameCHWSize image_size,int debug_mode);

    ~FaceMask();

    void pre_process(runtime_tensor& input_tensor, float* sparse_points);

    void inference();

    void post_process(FaceMaskInfo& result);

    void draw_results(cv::Mat& draw_frame, Bbox& bbox, FaceMaskInfo& result);

    void run(runtime_tensor& input_tensor,float* sparse_points, FaceMaskInfo& result);

protected:
    void svd22(const float a[4], float u[4], float s[2], float v[4]);
    
    void image_umeyama_128(float* src, float* dst);

    void get_affine_matrix(float* sparse_points);

    void softmax(vector<float>& input,vector<float>& output);

private:
    runtime_tensor model_input_tensor_;            
    FrameCHWSize image_size_;
    FrameCHWSize input_size_;

    float matrix_dst_[10];   
    float mask_thresh_;
    int debug_mode_;                     ///< Debug flag

    KPU *kpu = nullptr;                  ///< Pointer to KPU (Kendryte Processing Unit) instance
    AI2D *ai2d = nullptr;                ///< Pointer to AI2D preprocessing module
};
#endif