BOE_ZB_SOCKET/AIEngineModule/src/AI_Factory.cpp

/*
 * @Author: xiewenji 527774126@qq.com
 * @Date: 2025-09-03 10:33:48
 * @LastEditors: xiewenji 527774126@qq.com
 * @LastEditTime: 2025-09-14 20:28:10
 * @FilePath: /AI_SO_Test/AIEngineModule/src/AI_Factory.cpp
 * @Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
 */
#include "AI_Factory.h"

#include "AIModel_Impl.h"

#include <fstream>
#include <map>

// 模型路径配置文件（键值对格式，示例："Defect":/home/aidlux/BOE/UseModel_ZB/KK/Defect.engine）
static const std::string MODEL_PATH_CONFIG_FILE = "/home/aidlux/BOE/UseModel_ZB/KK.txt";

// 去掉字符串首尾的空白字符和引号
static std::string TrimStr(const std::string &str)
{
    size_t begin = str.find_first_not_of(" \t\r\n\"");
    if (begin == std::string::npos)
    {
        return "";
    }
    size_t end = str.find_last_not_of(" \t\r\n\"");
    return str.substr(begin, end - begin + 1);
}

// 读取模型路径键值对配置文件（每行格式：模型名:模型路径，支持 "模型名":模型路径 以及 # 注释行）
static std::map<std::string, std::string> ReadModelPathConfig(const std::string &strFilePath)
{
    std::map<std::string, std::string> mapConfig;
    std::ifstream ifs(strFilePath.c_str());
    if (!ifs.is_open())
    {
        printf("Open model config file failed, use default paths: %s\n", strFilePath.c_str());
        return mapConfig;
    }

    std::string strLine;
    while (std::getline(ifs, strLine))
    {
        // 跳过空行和注释行
        std::string strTrim = TrimStr(strLine);
        if (strTrim.empty() || strTrim[0] == '#')
        {
            continue;
        }
        size_t pos = strLine.find(':');
        if (pos == std::string::npos)
        {
            continue;
        }
        std::string strKey = TrimStr(strLine.substr(0, pos));
        std::string strValue = TrimStr(strLine.substr(pos + 1));
        if (!strKey.empty() && !strValue.empty())
        {
            mapConfig[strKey] = strValue;
        }
    }
    return mapConfig;
}

// 优先从配置中取模型路径，配置中没有则使用默认路径
static std::string GetModelPath(const std::map<std::string, std::string> &mapConfig,
                                const std::string &strKey, const std::string &strDefaultPath)
{
    std::map<std::string, std::string>::const_iterator it = mapConfig.find(strKey);
    if (it != mapConfig.end())
    {
        return it->second;
    }
    return strDefaultPath;
}

std::shared_ptr<AIModel_Base> AIModel_Base::GetInstance()
{
    std::shared_ptr<AIModel_Base> instance = std::make_shared<AIModel_Impl>();
    return instance;
}

AIFactory::AIFactory()
{
}

AIFactory::~AIFactory()
{
}
std::shared_ptr<AIFactory> AIFactory::GetInstance()
{
    static std::shared_ptr<AIFactory> instance = std::make_shared<AIFactory>();
    return instance;
}

int AIFactory::InitALLAIModle(GPU_Config gupconfig)
{
    if (gupconfig.GetNum() == 0)
    {
        printf("GPU Num Error \n");
        return 1;
    }

    // 读取模型路径配置文件，配置文件中有的路径优先，没有的用默认路径
    std::map<std::string, std::string> mapModelPaths = ReadModelPathConfig(MODEL_PATH_CONFIG_FILE);

    if (!Defect)
    {
        Defect = AIModel_Base::GetInstance();
        AIModel_Base::AIModelRun_Config edge_config;
        edge_config.gpuconfig.copy(gupconfig);
        edge_config.strPath = GetModelPath(mapModelPaths, "Defect", "/home/aidlux/BOE/UseModel_ZB/KK/Defect.engine");
        edge_config.strName = "Det";
        edge_config.inputType = AIModel_Base::Input_CHW;
        edge_config.Stream_num = 2;
        Defect->Init(edge_config);
    }
    if (!Defect_QueXi)
    {
        Defect_QueXi = AIModel_Base::GetInstance();
        AIModel_Base::AIModelRun_Config edge_config;
        edge_config.gpuconfig.copy(gupconfig);
        edge_config.strPath = GetModelPath(mapModelPaths, "QueXi", "/home/aidlux/BOE/UseModel_ZB/KK/QueXi.engine");
        edge_config.strName = "QueXi_Det";
        edge_config.inputType = AIModel_Base::Input_CHW;
        edge_config.Stream_num = 2;
        Defect_QueXi->Init(edge_config);
    }
    if (!Class)
    {
        Class = AIModel_Base::GetInstance();
        AIModel_Base::AIModelRun_Config jbl_config;
        jbl_config.gpuconfig.copy(gupconfig);
        jbl_config.strPath = GetModelPath(mapModelPaths, "Class", "/home/aidlux/BOE/UseModel_ZB/KK/Class.engine");
        jbl_config.inputType = AIModel_Base::Input_CHW;
        jbl_config.strName = "Class";
        jbl_config.IsClass = true;
        Class->Init(jbl_config);
    }
    if (!Pin_Loc)
    {
        Pin_Loc = AIModel_Base::GetInstance();
        AIModel_Base::AIModelRun_Config edge_config;
        edge_config.gpuconfig.copy(gupconfig);
        edge_config.strPath = GetModelPath(mapModelPaths, "PinLoc", "/home/aidlux/BOE/UseModel_ZB/KK/PinLoc.engine");
        edge_config.strName = "Pin_Loc";
        edge_config.inputType = AIModel_Base::Input_CHW;
        edge_config.Stream_num = 2;
        Pin_Loc->Init(edge_config);
    }
    if (!Align_Outer)
    {
        Align_Outer = AIModel_Base::GetInstance();
        AIModel_Base::AIModelRun_Config edge_config;
        edge_config.gpuconfig.copy(gupconfig);
        edge_config.strPath = GetModelPath(mapModelPaths, "AlignOuter", "/home/aidlux/BOE/UseModel_ZB/KK/AlignOuter.engine");
        edge_config.strName = "Align_Outer";
        edge_config.inputType = AIModel_Base::Input_CHW;
        edge_config.Stream_num = 2;
        Align_Outer->Init(edge_config);
    }
    if (!Align_Inner)
    {
        Align_Inner = AIModel_Base::GetInstance();
        AIModel_Base::AIModelRun_Config edge_config;
        edge_config.gpuconfig.copy(gupconfig);
        edge_config.strPath = GetModelPath(mapModelPaths, "AlignInner", "/home/aidlux/BOE/UseModel_ZB/KK/AlignInner.engine");
        edge_config.strName = "Align_Inner";
        edge_config.inputType = AIModel_Base::Input_CHW;
        edge_config.Stream_num = 2;
        Align_Inner->Init(edge_config);
    }
    m_bInitSucc = true;
    return 0;
}