更新 util/ModelUtils.py

util/ModelUtils.py

@@ -17,7 +17,7 @@ from util.PlotsUtils import get_label_arrays, get_label_array_dict
 from util.TorchUtils import select_device
 
 sys.path.extend(['..', '../AIlib2'])
-from AI import AI_process, AI_process_forest, get_postProcess_para, ocr_process, AI_process_N, AI_process_C
+from AI import AI_process, AI_process_forest, get_postProcess_para, ocr_process, AI_process_N, AI_process_C,AI_process_Ocr,AI_process_Crowd
 from stdc import stdcModel
 from segutils.segmodel import SegModel
 from models.experimental import attempt_load
@@ -241,7 +241,7 @@ def channel2_process(args):
         logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), request_id)
         raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0],
                                ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1])
-    
+
 def get_label_arraylist(*args):
     width, height, names, rainbows = args
     # line = int(round(0.002 * (height + width) / 2) + 1)
@@ -354,6 +354,73 @@ def im_process(args):
         raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0],
                                ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1])
 
+def immulti_process(args):
+    model_conf, frame, requestId = args
+    device, modelList, detpar = model_conf[1], model_conf[2], model_conf[3]
+    try:
+        # new_device = torch.device(device)
+        # img, padInfos = pre_process(frame, new_device)
+        # pred = model(img)
+        # boxes = post_process(pred, padInfos, device, conf_thres=pardet['conf_thres'],
+        #                      iou_thres=pardet['iou_thres'], nc=pardet['nc'])  # 后处理
+        return AI_process_Ocr([frame], modelList, device, detpar)
+    except ServiceException as s:
+        raise s
+    except Exception:
+        logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), requestId)
+        raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0],
+                               ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1])
+
+class CARPLATEModel:
+    __slots__ = "model_conf"
+    def __init__(self, device, allowedList=None, requestId=None, modeType=None, gpu_name=None, base_dir=None,
+                 env=None):
+        try:
+            logger.info("########################加载{}########################, requestId:{}", modeType.value[2],
+                        requestId)
+            par = modeType.value[4](str(device), gpu_name)
+            modelList=[ modelPar['model'](weights=modelPar['weight'],par=modelPar['par'])  for  modelPar in par['models'] ]
+            detpar = par['models'][0]['par']
+            # new_device = torch.device(par['device'])
+            # modelList=[ modelPar['model'](weights=modelPar['weight'],par=modelPar['par'])  for  modelPar in par['models'] ]
+            logger.info("########################加载 plate_yolov5s_v3.jit 成功 ########################, requestId:{}",
+                        requestId)
+            self.model_conf = (modeType, device, modelList, detpar, par['rainbows'])
+        except Exception:
+            logger.error("模型加载异常：{}, requestId:{}", format_exc(), requestId)
+            raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
+                                   ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
+
+
+class DENSECROWDCOUNTModel:
+    __slots__ = "model_conf"
+
+    def __init__(self, device, allowedList=None, requestId=None, modeType=None, gpu_name=None, base_dir=None, env=None):
+        try:
+            logger.info("########################加载{}########################, requestId:{}", modeType.value[2],
+                        requestId)
+            par = modeType.value[4](str(device), gpu_name)
+            rainbows = par["rainbows"]
+            models=[ modelPar['model'](weights=modelPar['weight'],par=modelPar['par'])  for  modelPar in par['models'] ]
+            postPar = par['models'][0]['par']
+            self.model_conf = (modeType, device, models[0], postPar, rainbows)
+        except Exception:
+            logger.error("模型加载异常：{}, requestId:{}", format_exc(), requestId)
+            raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
+                                   ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
+
+def cc_process(args):
+    model_conf, frame, requestId = args
+    device, model, postPar = model_conf[1], model_conf[2], model_conf[3]
+    try:
+        return AI_process_Crowd([frame], model, device, postPar)
+    except ServiceException as s:
+        raise s
+    except Exception:
+        logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), requestId)
+        raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0],
+                               ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1])
+
 
 # 百度AI图片识别模型
 class BaiduAiImageModel:
@@ -470,7 +537,7 @@ MODEL_CONFIG = {
         lambda x, y, z: one_label(x, y, z),
         lambda x: detSeg_demo2(x)
     ),
-    
+
     # 加载交通模型
     ModelType.TRAFFIC_FARM_MODEL.value[1]: (
         lambda x, y, r, t, z, h: OneModel(x, y, r, ModelType.TRAFFIC_FARM_MODEL, t, z, h),
@@ -613,22 +680,22 @@ MODEL_CONFIG = {
         lambda x, y, z: one_label(x, y, z),
         lambda x: model_process(x)
     ),
-    # 加载智慧工地模型    
+    # 加载智慧工地模型
     ModelType.SMARTSITE_MODEL.value[1]: (
         lambda x, y, r, t, z, h: cityManagementModel(x, y, r, ModelType.SMARTSITE_MODEL, t, z, h),
         ModelType.SMARTSITE_MODEL,
         lambda x, y, z: one_label(x, y, z),
         lambda x: detSeg_demo2(x)
     ),
-    
+
-    # 加载垃圾模型    
+    # 加载垃圾模型
     ModelType.RUBBISH_MODEL.value[1]: (
         lambda x, y, r, t, z, h: cityManagementModel(x, y, r, ModelType.RUBBISH_MODEL, t, z, h),
         ModelType.RUBBISH_MODEL,
         lambda x, y, z: one_label(x, y, z),
         lambda x: detSeg_demo2(x)
     ),
-    
+
     # 加载烟花模型    
     ModelType.FIREWORK_MODEL.value[1]: (
         lambda x, y, r, t, z, h: cityManagementModel(x, y, r, ModelType.FIREWORK_MODEL, t, z, h),
@@ -657,6 +724,13 @@ MODEL_CONFIG = {
         lambda x, y, z: one_label(x, y, z),
         lambda x: detSeg_demo2(x)
     ),
+    # 加载自研车牌检测模型
+    ModelType.CITY_CARPLATE_MODEL.value[1]: (
+        lambda x, y, r, t, z, h: CARPLATEModel(x, y, r, ModelType.CITY_CARPLATE_MODEL, t, z, h),
+        ModelType.CITY_CARPLATE_MODEL,
+        None,
+        lambda x: immulti_process(x)
+    ),
     # 加载红外行人检测模型
     ModelType.CITY_INFRAREDPERSON_MODEL.value[1]: (
         lambda x, y, r, t, z, h: cityManagementModel(x, y, r, ModelType.CITY_INFRAREDPERSON_MODEL, t, z, h),
@@ -670,5 +744,12 @@ MODEL_CONFIG = {
         ModelType.CITY_NIGHTFIRESMOKE_MODEL,
         lambda x, y, z: one_label(x, y, z),
         lambda x: detSeg_demo2(x)
-    ),    
+    ),
+    # 加载密集人群计数检测模型
+    ModelType.CITY_DENSECROWDCOUNT_MODEL.value[1]: (
+        lambda x, y, r, t, z, h: DENSECROWDCOUNTModel(x, y, r, ModelType.CITY_DENSECROWDCOUNT_MODEL, t, z, h),
+        ModelType.CITY_DENSECROWDCOUNT_MODEL,
+        None,
+        lambda x: cc_process(x)
+    ),
 }