diff --git a/AIEngineModule/src/AI_Factory.cpp b/AIEngineModule/src/AI_Factory.cpp index 488e313..86670f5 100644 --- a/AIEngineModule/src/AI_Factory.cpp +++ b/AIEngineModule/src/AI_Factory.cpp @@ -10,6 +10,71 @@ #include "AIModel_Impl.h" +#include +#include + +// 模型路径配置文件(键值对格式,示例:"01DefectLight":/home/aidlux/BOE/UseModel_ZB/BM/01DefectLight.engine) +static const std::string MODEL_PATH_CONFIG_FILE = "/home/aidlux/BOE/UseModel_ZB/BM.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 ReadModelPathConfig(const std::string &strFilePath) +{ + std::map 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 &mapConfig, + const std::string &strKey, const std::string &strDefaultPath) +{ + std::map::const_iterator it = mapConfig.find(strKey); + if (it != mapConfig.end()) + { + return it->second; + } + return strDefaultPath; +} + std::shared_ptr AIModel_Base::GetInstance() { std::shared_ptr instance = std::make_shared(); @@ -37,12 +102,15 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig) return 1; } + // 读取模型路径配置文件,配置文件中有的路径优先,没有的用默认路径 + std::map 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 = "/home/aidlux/BOE/UseModel_ZB/BM/01DefectLight.engine"; + edge_config.strPath = GetModelPath(mapModelPaths, "01DefectLight", "/home/aidlux/BOE/UseModel_ZB/BM/01DefectLight.engine"); edge_config.strName = "DefectLight"; edge_config.inputType = AIModel_Base::Input_CHW; edge_config.Stream_num = 4; @@ -53,7 +121,7 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig) Defect_Dark = AIModel_Base::GetInstance(); AIModel_Base::AIModelRun_Config edge_config; edge_config.gpuconfig.copy(gupconfig); - edge_config.strPath = "/home/aidlux/BOE/UseModel_ZB/BM/05DefectDark.engine"; + edge_config.strPath = GetModelPath(mapModelPaths, "05DefectDark", "/home/aidlux/BOE/UseModel_ZB/BM/05DefectDark.engine"); edge_config.strName = "DefectDark"; edge_config.inputType = AIModel_Base::Input_CHW; edge_config.Stream_num = 4; @@ -64,7 +132,7 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig) Class = AIModel_Base::GetInstance(); AIModel_Base::AIModelRun_Config jbl_config; jbl_config.gpuconfig.copy(gupconfig); - jbl_config.strPath = "/home/aidlux/BOE/UseModel_ZB/BM/Class.engine"; + jbl_config.strPath = GetModelPath(mapModelPaths, "Class", "/home/aidlux/BOE/UseModel_ZB/BM/Class.engine"); jbl_config.inputType = AIModel_Base::Input_CHW; jbl_config.strName = "Class"; jbl_config.IsClass = true; @@ -75,7 +143,7 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig) Tag_Loc = AIModel_Base::GetInstance(); AIModel_Base::AIModelRun_Config edge_config; edge_config.gpuconfig.copy(gupconfig); - edge_config.strPath = "/home/aidlux/BOE/UseModel_ZB/BM/02BadLabel.engine"; + edge_config.strPath = GetModelPath(mapModelPaths, "02BadLabel", "/home/aidlux/BOE/UseModel_ZB/BM/02BadLabel.engine"); edge_config.strName = "Tag_Loc"; edge_config.inputType = AIModel_Base::Input_CHW; edge_config.Stream_num = 4; @@ -86,7 +154,7 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig) Align_Outer = AIModel_Base::GetInstance(); AIModel_Base::AIModelRun_Config edge_config; edge_config.gpuconfig.copy(gupconfig); - edge_config.strPath = "/home/aidlux/BOE/UseModel_ZB/BM/04ProductLoc.engine"; + edge_config.strPath = GetModelPath(mapModelPaths, "04ProductLoc", "/home/aidlux/BOE/UseModel_ZB/BM/04ProductLoc.engine"); edge_config.strName = "Align_Outer"; edge_config.inputType = AIModel_Base::Input_CHW; edge_config.Stream_num = 2; @@ -97,7 +165,7 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig) Align_Inner = AIModel_Base::GetInstance(); AIModel_Base::AIModelRun_Config edge_config; edge_config.gpuconfig.copy(gupconfig); - edge_config.strPath = "/home/aidlux/BOE/UseModel_ZB/BM/03ClampLoc.engine"; + edge_config.strPath = GetModelPath(mapModelPaths, "03ClampLoc", "/home/aidlux/BOE/UseModel_ZB/BM/03ClampLoc.engine"); edge_config.strName = "Align_Inner"; edge_config.inputType = AIModel_Base::Input_CHW; edge_config.Stream_num = 2;