Embedded_AI/npu/axera/include/AxeraModel.h

#pragma once

#include <string>
#include <chrono>
#include "nlohmann/json.hpp"
#include "Common.h"
#include "AIStatus.h"

#include "file.hpp"
#include "io.hpp"
#include "ax_sys_api.h"
#include "joint_adv.h"
#include "cv/cv.hpp"
#include "cv/utils.hpp"
#include "Detection.hpp"

class ax_crop_resize_nv12
{
public:
    ax_crop_resize_nv12() = default;
    int init(int input_h, int input_w, int model_h, int model_w, int ModelConfigype);
	int deinit();
    int run_crop_resize_nv12(void *input_data, void *output_data, int &output_data_size);
    // int run_crop_resize_nv12_new(void *input_data);
	// AX_NPU_CV_Image* m_output_image;
private:
    AX_NPU_CV_Image* m_input_image;
    AX_NPU_CV_Image* m_output_image;
    AX_NPU_CV_Box* m_box;

    int m_model_type = AX_NPU_MODEL_TYPE_1_1_1;
};

class AxeraModel
{
public:
    /**
	 * @brief 预处理结构体:用于表示图像预处理的状态和配置。
	 */
	struct PreprocessConfig
	{
		// preprocess resource
		void *dstVirtualAddress;
		int dstImageSize;
		ax_crop_resize_nv12 axcv_imageconvert;

		// net wh
		int netWidth;
		int netHeight;
		// src wh
		int srcWidth;
		int srcHeight;
		float scaleInfo;
    };

    /**
	 * @brief 模型结构体:用于表示神经网络模型的状态和配置。
	 */
    struct ModelConfig
	{
		// runtime handle
		AX_JOINT_HANDLE joint_handle;
		AX_JOINT_SDK_ATTR_T joint_attr;

		// context
		AX_JOINT_EXECUTION_CONTEXT joint_ctx;
        AX_JOINT_EXECUTION_CONTEXT_SETTING_T joint_ctx_settings;

		// prepare_io resurce
		AX_JOINT_IO_T joint_io_arr;
        AX_JOINT_IO_SETTING_T joint_io_setting;

		const AX_JOINT_IO_INFO_T *io_info;

		// result
		int netWidth, netHeight;							///< 网络的宽高
		std::vector<detection::Box> objects;
		std::vector<float_t *> outputBlobsPtr; ///< 输出 blob 的张量指针（request用）
		std::vector<int> outputBlobsPtrSize; ///< 输出 blob 的张量指针（request用）
		std::vector<std::array<size_t, 4>>  outputDataDims;
		unsigned int strides[5] = {0}; ///< 步长数组

		std::string modelType;						///< 模型类型
		std::vector<std::string> classNames;		///< 类别名称数组
    };

	struct PostprocessConfig
	{
		std::string name;
		unsigned int yoloAnchor[30] = {0};			///< yolo_anchor数组
	};


    /**
	 * @brief 初始化预处理资源
	 *
	 * 根据提供的配置信息初始化预处理资源。
	 *
	 * @param cfg 预处理配置的 JSON 对象
	 * @param preprocessConfig 预处理状态结构体指针
	 * @return 表示预处理资源初始化成功或失败的错误代码。
	 */
	int initializePreprocessResource(const nlohmann::json_abi_v3_11_2::json cfg, PreprocessConfig *preprocessConfig);
	/**
	 * @brief 初始化模型资源
	 *
	 * 根据提供的配置信息初始化模型资源。
	 *
	 * @param cfg 模型配置的 JSON 对象
	 * @param modelConfig 模型状态结构体指针
	 * @return 表示模型资源初始化成功或失败的错误代码。
	 */
	int initializeModelResource(const nlohmann::json_abi_v3_11_2::json cfg, ModelConfig *modelConfig);
	/**
	 * @brief 预处理图像
	 *
	 * 使用提供的预处理状态和图像帧状态执行图像预处理。
	 *
	 * @param preprocessConfig 预处理状态结构体指针
	 * @param frame_st 图像帧状态结构体指针
	 * @return 表示预处理图像操作成功或失败的错误代码。
	 */
	int preprocessImage(PreprocessConfig *preprocessConfig, frame_t *frame_st);
	/**
	 * @brief 模型推理
	 *
	 * 使用提供的预处理状态和模型状态执行推理。
	 *
	 * @param preprocessConfig 预处理状态结构体指针
	 * @param modelConfig 模型状态结构体指针
     * @param infer_result 存储模型结果的指针
	 * @return 表示模型推理成功或失败的错误代码。
	 */
	int inferModel(unsigned char *preProcessedImg,int preProcessedImgSize, ModelConfig *modelConfig);
	/**
	 * @brief 释放预处理资源
	 *
	 * 使用提供的预处理状态释放相关资源。
	 *
	 * @param preprocessConfig 预处理状态结构体指针
	 * @return 表示释放预处理资源成功或失败的错误代码。
	 */
	int postprocess(std::vector<BoundingBox> *boundingBoxes,
					std::vector<float *> outputBlobsPtr,
					std::vector<std::array<size_t, 4>> outputDataDims,
					unsigned int strides[5],
					unsigned int yoloAnchor[30],
					std::vector<std::string> classNames,
					std::unordered_map<std::string, float> classAttributes,
					float scaleInfo, int srcWidth,int srcHeight);
	int releasePreprocessResources(PreprocessConfig *preprocessConfig);
	/**
	 * @brief 释放模型资源
	 *
	 * 使用提供的模型状态释放相关资源。
	 *
	 * @param modelConfig 模型状态结构体指针
	 * @return 表示释放模型资源成功或失败的错误代码。
	 */
	int releaseModelResources(ModelConfig *modelConfig);
};