gitadmin
/
tuoheng_algN
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

model_process:支持按类置信度筛选类别

This commit is contained in:
th 2025-07-31 19:33:00 +08:00
parent b46182d68e
commit 2d626929f7
2 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
enums/ModelTypeEnum.py
View File

@ -162,6 +162,7 @@ class ModelType(Enum):
"classes": 10, "classes": 10,
"rainbows": COLOR "rainbows": COLOR
}, },
'score_byClass':{11:0.75,12:0.75},
'allowedList':[0,1,2,3,4,5,6,7,8,9,10,11,12,16,17,18,19,20,21,22], 'allowedList':[0,1,2,3,4,5,6,7,8,9,10,11,12,16,17,18,19,20,21,22],
'Detweights': "../weights/trt/AIlib2/highWay2/yolov5_%s_fp16.engine" % gpuName, 'Detweights': "../weights/trt/AIlib2/highWay2/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/highWay2/stdc_360X640_%s_fp16.engine' % gpuName 'Segweights': '../weights/trt/AIlib2/highWay2/stdc_360X640_%s_fp16.engine' % gpuName

5
util/ModelUtils.py
View File

@ -35,6 +35,7 @@ class OneModel:
__slots__ = "model_conf" __slots__ = "model_conf"
def __init__(self, device, allowedList=None, requestId=None, modeType=None, gpu_name=None, base_dir=None, env=None): def __init__(self, device, allowedList=None, requestId=None, modeType=None, gpu_name=None, base_dir=None, env=None):
s = time.time()
try: try:
logger.info("########################加载{}########################, requestId:{}", modeType.value[2], logger.info("########################加载{}########################, requestId:{}", modeType.value[2],
requestId) requestId)
@ -71,7 +72,8 @@ class OneModel:
'allowedList': par['allowedList'] if modeType.value[0] == '3' else [], 'allowedList': par['allowedList'] if modeType.value[0] == '3' else [],
'segRegionCnt': par['segRegionCnt'], 'segRegionCnt': par['segRegionCnt'],
'trtFlag_det': par['trtFlag_det'], 'trtFlag_det': par['trtFlag_det'],
'trtFlag_seg': par['trtFlag_seg'] 'trtFlag_seg': par['trtFlag_seg'],
'score_byClass':par['score_byClass'] if 'score_byClass' in par.keys() else None
} }
model_param = { model_param = {
"model": model, "model": model,
@ -86,6 +88,7 @@ class OneModel:
logger.error("模型加载异常：{}, requestId:{}", format_exc(), requestId) logger.error("模型加载异常：{}, requestId:{}", format_exc(), requestId)
raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0], raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
ExceptionType.MODEL_LOADING_EXCEPTION.value[1]) ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
logger.info("模型初始化时间：{}, requestId:{}", time.time() - s, requestId)
# 纯分类模型 # 纯分类模型
class cityManagementModel: class cityManagementModel:
__slots__ = "model_conf" __slots__ = "model_conf"