feat 开放模型路径配置文件

dev_lsy
liusiyang 3 weeks ago
parent c7b20f2c75
commit f8999aeb34

@ -10,6 +10,71 @@
#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>();
@ -37,12 +102,15 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig)
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 = "/home/aidlux/BOE/UseModel_ZB/KK/Defect.engine";
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;
@ -53,7 +121,7 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig)
Defect_QueXi = AIModel_Base::GetInstance();
AIModel_Base::AIModelRun_Config edge_config;
edge_config.gpuconfig.copy(gupconfig);
edge_config.strPath = "/home/aidlux/BOE/UseModel_ZB/KK/QueXi.engine";
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;
@ -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/KK/Class.engine";
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;
@ -75,7 +143,7 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig)
Pin_Loc = AIModel_Base::GetInstance();
AIModel_Base::AIModelRun_Config edge_config;
edge_config.gpuconfig.copy(gupconfig);
edge_config.strPath = "/home/aidlux/BOE/UseModel_ZB/KK/PinLoc.engine";
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;
@ -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/KK/AlignOuter.engine";
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;
@ -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/KK/AlignInner.engine";
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;

Loading…
Cancel
Save