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

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

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"