Compare commits

...

3 Commits

@ -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;

@ -13,6 +13,179 @@
#include <map>
#include <algorithm>
#include "AICommonDefine.h"
#include <cstdlib>
#include <fstream>
// 分类映射配置文件(每行:label索引:缺陷名称;冒号后名称为该缺陷类型名称,行的顺序即缺陷类型索引)
static const std::string CLS_LABEL_CONFIG_FILE = "/home/aidlux/BOE/UseModel_ZB/KK_cls.txt";
// 去掉字符串首尾的空白字符和引号
static std::string ClsTrimStr(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);
}
// 组装程序内默认缺陷名称列表(AI_CLass_QX_NAME_Names)
static std::vector<std::string> BuildDefaultAIQXNames()
{
std::vector<std::string> vecNames;
for (int i = 0; i < AI_CLass_QX_NAME_count; i++)
{
vecNames.push_back(AI_CLass_QX_NAME_Names[i]);
}
return vecNames;
}
// 分类映射配置:缺陷名称列表(索引=行序) + 分类label->缺陷名称 映射
struct ClsLabelConfig
{
std::vector<std::string> vecNames; // 缺陷名称列表
std::map<int, std::string> mapLabelName; // 分类 label -> 缺陷名称
};
// 读取分类映射配置(每行:label索引:缺陷名称;冒号后名称按行序构成缺陷名称列表;# 后内容忽略)
static ClsLabelConfig ReadClsLabelConfig(const std::string &strFilePath)
{
ClsLabelConfig config;
std::ifstream ifs(strFilePath.c_str());
if (!ifs.is_open())
{
printf("Open cls label config failed, use default mapping: %s\n", strFilePath.c_str());
return config;
}
std::string strLine;
while (std::getline(ifs, strLine))
{
// # 后内容(含注释)忽略
size_t nPosComment = strLine.find('#');
if (nPosComment != std::string::npos)
{
strLine = strLine.substr(0, nPosComment);
}
size_t nPos = strLine.find(':');
if (nPos == std::string::npos)
{
// 无冒号:非标准行,跳过并提示(仅支持 label索引:缺陷名称 格式)
std::string strTrim = ClsTrimStr(strLine);
if (!strTrim.empty())
{
printf("Cls label config skip invalid line: %s\n", strTrim.c_str());
}
continue;
}
// label索引:缺陷名称(冒号后名称按行序构成缺陷名称列表)
std::string strKey = ClsTrimStr(strLine.substr(0, nPos));
std::string strValue = ClsTrimStr(strLine.substr(nPos + 1));
if (!strKey.empty() && !strValue.empty())
{
config.mapLabelName[atoi(strKey.c_str())] = strValue;
config.vecNames.push_back(strValue); // 行序即缺陷类型索引
}
}
return config;
}
// 获取分类映射配置(只读一次)
static const ClsLabelConfig &GetClsLabelConfig()
{
static const ClsLabelConfig config = ReadClsLabelConfig(CLS_LABEL_CONFIG_FILE);
return config;
}
// 获取缺陷名称列表:优先取配置文件的名称列表,读不到用程序内默认列表
static const std::vector<std::string> &GetAIQXNameList()
{
static const std::vector<std::string> vecDefaultNames = BuildDefaultAIQXNames();
const ClsLabelConfig &config = GetClsLabelConfig();
if (!config.vecNames.empty())
{
return config.vecNames;
}
return vecDefaultNames;
}
// 缺陷名称 -> AI_CLass 枚举(找不到返回 count 表示无效)
static int ClsNameToAItype(const std::string &strName)
{
const std::vector<std::string> &vecNames = GetAIQXNameList();
for (size_t i = 0; i < vecNames.size(); i++)
{
if (strName == vecNames[i])
{
return (int)i;
}
}
return AI_CLass_QX_NAME_count;
}
// 解析 label->缺陷名称 映射为 label->AI_CLass枚举(只读一次;无效名称仅提示一次)
static std::map<int, int> BuildLabelTypeMap()
{
const ClsLabelConfig &config = GetClsLabelConfig();
std::map<int, int> mapLabelType;
for (std::map<int, std::string>::const_iterator it = config.mapLabelName.begin(); it != config.mapLabelName.end(); ++it)
{
int nType = ClsNameToAItype(it->second);
if (nType < AI_CLass_QX_NAME_count)
{
mapLabelType[it->first] = nType;
}
else
{
printf("Cls label config unknown name: %s\n", it->second.c_str());
}
}
return mapLabelType;
}
// 获取分类 label 对应的 AI 缺陷类型:优先配置文件(只读一次),读不到用默认顺序
static int GetClsTypeByLabel(int nLabel)
{
static const std::map<int, int> mapLabelType = BuildLabelTypeMap();
std::map<int, int>::const_iterator it = mapLabelType.find(nLabel);
if (it != mapLabelType.end())
{
return it->second;
}
// 映射文件未读到、或该 label 未在文件中配置时,按程序内默认顺序兜底(避免文件缺失/不全导致异常)
int cls_num = AI_CLass_QX_NAME_count;
switch (nLabel)
{
case 0:
cls_num = AI_CLass_QX_NAME_PIN_PoSun;
break;
case 1:
cls_num = AI_CLass_QX_NAME_PIN_QueShi;
break;
case 2:
cls_num = AI_CLass_QX_NAME_BTB_DuanLie;
break;
case 3:
cls_num = AI_CLass_QX_NAME_LianXi;
break;
case 4:
cls_num = AI_CLass_QX_NAME_QueXi;
break;
case 5:
cls_num = AI_CLass_QX_NAME_YiWu;
break;
case 6:
cls_num = AI_CLass_QX_NAME_C_BianXing;
break;
default:
break;
}
return cls_num;
}
// 用于排序轮廓的比较函数
static bool compareContourAreas(const vector<Point> &contour1, const vector<Point> &contour2)
{
@ -183,7 +356,7 @@ int ImgCheckAnalysisy::GetStatus()
std::string ImgCheckAnalysisy::GetVersion()
{
return std::string("BOE_1.2.0");
return std::string("BOE_1.3.0");
}
std::string ImgCheckAnalysisy::GetErrorInfo()
@ -1654,42 +1827,9 @@ int ImgCheckAnalysisy::AI_QX_Class_Thread()
{
int temClass = result->cls_label;
int cls_num = 0;
switch (temClass)
{
case 0:
cls_num = AI_CLass_QX_NAME_PIN_PoSun;
break;
case 1:
cls_num = AI_CLass_QX_NAME_PIN_QueShi;
break;
case 2:
cls_num = AI_CLass_QX_NAME_BTB_DuanLie;
break;
case 3:
cls_num = AI_CLass_QX_NAME_LianXi;
break;
case 4:
cls_num = AI_CLass_QX_NAME_QueXi;
break;
case 5:
cls_num = AI_CLass_QX_NAME_YiWu;
break;
case 6:
cls_num = AI_CLass_QX_NAME_C_BianXing;
break;
default:
cls_num = AI_CLass_QX_NAME_count; // 无效类别, 不参与缺陷分析
break;
}
// 分类 label 映射:优先读映射配置文件(KK_cls.txt),读不到用默认顺序
int cls_num = GetClsTypeByLabel(temClass);
// std::string strclassName = AI_CLass_QX_NAME_Names[cls_num];
// printf("AI Class num %d = %s %f\n", cls_num, strclassName.c_str(), result->cls_score);

Loading…
Cancel
Save