feat 开放模型路径配置文件

dev_lsy
liusiyang 3 weeks ago
parent 1ce1c502e8
commit f21d3d95df

@ -10,6 +10,71 @@
#include "AIModel_Impl.h" #include "AIModel_Impl.h"
#include <fstream>
#include <map>
// 模型路径配置文件(键值对格式,示例:"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<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> AIModel_Base::GetInstance()
{ {
std::shared_ptr<AIModel_Base> instance = std::make_shared<AIModel_Impl>(); std::shared_ptr<AIModel_Base> instance = std::make_shared<AIModel_Impl>();
@ -37,12 +102,15 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig)
return 1; return 1;
} }
// 读取模型路径配置文件,配置文件中有的路径优先,没有的用默认路径
std::map<std::string, std::string> mapModelPaths = ReadModelPathConfig(MODEL_PATH_CONFIG_FILE);
if (!Defect) if (!Defect)
{ {
Defect = AIModel_Base::GetInstance(); Defect = AIModel_Base::GetInstance();
AIModel_Base::AIModelRun_Config edge_config; AIModel_Base::AIModelRun_Config edge_config;
edge_config.gpuconfig.copy(gupconfig); 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.strName = "DefectLight";
edge_config.inputType = AIModel_Base::Input_CHW; edge_config.inputType = AIModel_Base::Input_CHW;
edge_config.Stream_num = 4; edge_config.Stream_num = 4;
@ -53,7 +121,7 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig)
Defect_Dark = AIModel_Base::GetInstance(); Defect_Dark = AIModel_Base::GetInstance();
AIModel_Base::AIModelRun_Config edge_config; AIModel_Base::AIModelRun_Config edge_config;
edge_config.gpuconfig.copy(gupconfig); 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.strName = "DefectDark";
edge_config.inputType = AIModel_Base::Input_CHW; edge_config.inputType = AIModel_Base::Input_CHW;
edge_config.Stream_num = 4; edge_config.Stream_num = 4;
@ -64,7 +132,7 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig)
Class = AIModel_Base::GetInstance(); Class = AIModel_Base::GetInstance();
AIModel_Base::AIModelRun_Config jbl_config; AIModel_Base::AIModelRun_Config jbl_config;
jbl_config.gpuconfig.copy(gupconfig); 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.inputType = AIModel_Base::Input_CHW;
jbl_config.strName = "Class"; jbl_config.strName = "Class";
jbl_config.IsClass = true; jbl_config.IsClass = true;
@ -75,7 +143,7 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig)
Tag_Loc = AIModel_Base::GetInstance(); Tag_Loc = AIModel_Base::GetInstance();
AIModel_Base::AIModelRun_Config edge_config; AIModel_Base::AIModelRun_Config edge_config;
edge_config.gpuconfig.copy(gupconfig); 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.strName = "Tag_Loc";
edge_config.inputType = AIModel_Base::Input_CHW; edge_config.inputType = AIModel_Base::Input_CHW;
edge_config.Stream_num = 4; edge_config.Stream_num = 4;
@ -86,7 +154,7 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig)
Align_Outer = AIModel_Base::GetInstance(); Align_Outer = AIModel_Base::GetInstance();
AIModel_Base::AIModelRun_Config edge_config; AIModel_Base::AIModelRun_Config edge_config;
edge_config.gpuconfig.copy(gupconfig); 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.strName = "Align_Outer";
edge_config.inputType = AIModel_Base::Input_CHW; edge_config.inputType = AIModel_Base::Input_CHW;
edge_config.Stream_num = 2; edge_config.Stream_num = 2;
@ -97,7 +165,7 @@ int AIFactory::InitALLAIModle(GPU_Config gupconfig)
Align_Inner = AIModel_Base::GetInstance(); Align_Inner = AIModel_Base::GetInstance();
AIModel_Base::AIModelRun_Config edge_config; AIModel_Base::AIModelRun_Config edge_config;
edge_config.gpuconfig.copy(gupconfig); 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.strName = "Align_Inner";
edge_config.inputType = AIModel_Base::Input_CHW; edge_config.inputType = AIModel_Base::Input_CHW;
edge_config.Stream_num = 2; edge_config.Stream_num = 2;

Loading…
Cancel
Save