diff --git a/enums/ModelTypeEnum.py b/enums/ModelTypeEnum.py index d2adeba..c80852c 100644 --- a/enums/ModelTypeEnum.py +++ b/enums/ModelTypeEnum.py @@ -162,6 +162,7 @@ class ModelType(Enum): "classes": 10, "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], 'Detweights': "../weights/trt/AIlib2/highWay2/yolov5_%s_fp16.engine" % gpuName, 'Segweights': '../weights/trt/AIlib2/highWay2/stdc_360X640_%s_fp16.engine' % gpuName diff --git a/util/ModelUtils.py b/util/ModelUtils.py index 4dd2f7e..eefefdf 100644 --- a/util/ModelUtils.py +++ b/util/ModelUtils.py @@ -35,6 +35,7 @@ class OneModel: __slots__ = "model_conf" def __init__(self, device, allowedList=None, requestId=None, modeType=None, gpu_name=None, base_dir=None, env=None): + s = time.time() try: logger.info("########################加载{}########################, requestId:{}", modeType.value[2], requestId) @@ -71,7 +72,8 @@ class OneModel: 'allowedList': par['allowedList'] if modeType.value[0] == '3' else [], 'segRegionCnt': par['segRegionCnt'], '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": model, @@ -86,6 +88,7 @@ class OneModel: logger.error("模型加载异常:{}, requestId:{}", format_exc(), requestId) raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0], ExceptionType.MODEL_LOADING_EXCEPTION.value[1]) + logger.info("模型初始化时间:{}, requestId:{}", time.time() - s, requestId) # 纯分类模型 class cityManagementModel: __slots__ = "model_conf"