diff --git a/enums/ModelTypeEnum.py b/enums/ModelTypeEnum.py index 6f619f7..1f89471 100644 --- a/enums/ModelTypeEnum.py +++ b/enums/ModelTypeEnum.py @@ -9,12 +9,10 @@ from DMPRUtils.jointUtil import dmpr_yolo from segutils.segmodel import SegModel from utilsK.queRiver import riverDetSegMixProcess from utilsK.crowdGather import gather_post_process -from segutils.trafficUtils import tracfficAccidentMixFunction +from segutils.trafficUtils import tracfficAccidentMixFunction,mixTraffic_postprocess from utilsK.drownUtils import mixDrowing_water_postprocess from utilsK.noParkingUtils import mixNoParking_road_postprocess from utilsK.illParkingUtils import illParking_postprocess -from utilsK.spillUtils import mixSpillage_postprocess -from utilsK.cthcUtils import mixCthc_postprocess from utilsK.pannelpostUtils import pannel_post_process from stdc import stdcModel from yolov5 import yolov5Model @@ -119,7 +117,6 @@ class ModelType(Enum): }) - TRAFFIC_FARM_MODEL = ("3", "003", "交通模型", 'highWay2', lambda device, gpuName: { @@ -164,7 +161,7 @@ class ModelType(Enum): "classes": 10, "rainbows": COLOR }, - '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, 'Segweights': '../weights/trt/AIlib2/highWay2/stdc_360X640_%s_fp16.engine' % gpuName }) @@ -359,9 +356,7 @@ class ModelType(Enum): "classes": 5, "rainbows": COLOR }, - # "../weights/pth/AIlib2/%s/yolov5.pt" % modeType.value[3] 'Detweights': "../weights/trt/AIlib2/river2/yolov5_%s_fp16.engine" % gpuName, - # '../weights/pth/AIlib2/%s/stdc_360X640.pth' % modeType.value[3] 'Segweights': '../weights/trt/AIlib2/river2/stdc_360X640_%s_fp16.engine' % gpuName }) @@ -378,11 +373,10 @@ class ModelType(Enum): }, 'models':[ { - 'weight':'../weights/pth/AIlib2/cityMangement3/yolov5.pt', - #'weight':'../weights/trt/AIlib2/cityMangement3/yolov5_%s_fp16.engine'%(gpuName), + 'weight':'../weights/trt/AIlib2/cityMangement3/yolov5_%s_fp16.engine'%(gpuName), 'name':'yolov5', 'model':yolov5Model, - 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.4,"2":0.5,"3":0.5 } } + 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':True,'trtFlag_seg':True, "score_byClass":{"0":0.8,"1":0.4,"2":0.5,"3":0.5 } } }, { 'weight':'../weights/pth/AIlib2/cityMangement3/dmpr.pth', @@ -394,9 +388,9 @@ class ModelType(Enum): 'name':'dmpr' }, { - 'weight':'../weights/pth/AIlib2/cityMangement3/stdc_360X640.pth', + 'weight':'../weights/trt/AIlib2/cityMangement3/stdc_360X640_%s_fp16.engine'%(gpuName), 'par':{ - 'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'predResize':True,'numpy':False, 'RGB_convert_first':True,'seg_nclass':2},###分割模型预处理参数 + 'modelSize':(640,360),'mean':(0.485, 0.456, 0.406),'std' :(0.229, 0.224, 0.225),'predResize':True,'numpy':False, 'RGB_convert_first':True,'seg_nclass':3},###分割模型预处理参数 'model':stdcModel, 'name':'stdc' } @@ -441,9 +435,7 @@ class ModelType(Enum): "classes": 9, "rainbows": COLOR }, - # "../weights/pth/AIlib2/%s/yolov5.pt" % modeType.value[3] 'Detweights': "../weights/trt/AIlib2/drowning/yolov5_%s_fp16.engine" % gpuName, - # '../weights/pth/AIlib2/%s/stdc_360X640.pth' % modeType.value[3] 'Segweights': '../weights/trt/AIlib2/drowning/stdc_360X640_%s_fp16.engine' % gpuName }) @@ -512,7 +504,7 @@ class ModelType(Enum): CITYROAD_MODEL = ("20", "020", "城市公路模型", 'cityRoad', lambda device, gpuName: { 'device': device, - 'labelnames': ["护栏", "交通标志", "非交通标志", "施工", "施工"], + 'labelnames': ["护栏", "交通标志", "非交通标志", "施工锥桶", "施工水马"], 'trtFlag_seg': False, 'trtFlag_det': True, 'slopeIndex': [], @@ -572,14 +564,12 @@ class ModelType(Enum): }}, 'models':[ { - #'weight':'../weights/pth/AIlib2/channel2/yolov5.pt', 'weight':'../weights/trt/AIlib2/channel2/yolov5_%s_fp16.engine'%(gpuName), 'name':'yolov5', 'model':yolov5Model, 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.1,'iou_thres':0.45,'allowedList':list(range(20)),'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.7,"1":0.7,"2":0.8,"3":0.6} } }, - { - # 'weight' : '../weights/trt/AIlib2/ocr2/crnn_ch_4090_fp16_192X32.engine', + { 'weight' : '../weights/pth/AIlib2/ocr2/crnn_ch.pth', 'name':'ocr', 'model':ocrModel, @@ -638,9 +628,7 @@ class ModelType(Enum): "classes": 5, "rainbows": COLOR }, - # "../weights/pth/AIlib2/%s/yolov5.pt" % modeType.value[3] 'Detweights': "../weights/trt/AIlib2/riverT/yolov5_%s_fp16.engine" % gpuName, - # '../weights/pth/AIlib2/%s/stdc_360X640.pth' % modeType.value[3] 'Segweights': '../weights/trt/AIlib2/riverT/stdc_360X640_%s_fp16.engine' % gpuName }) @@ -784,26 +772,24 @@ class ModelType(Enum): 'labelnames': ["抛洒物","车辆"], 'trtFlag_seg': True, 'trtFlag_det': True, - 'seg_nclass': 2, + 'seg_nclass': 3, 'segRegionCnt': 2, 'segPar': { - #'modelSize': (640, 360), - 'modelSize': (1920, 1080), + 'modelSize': (640, 360), 'mean': (0.485, 0.456, 0.406), 'std': (0.229, 0.224, 0.225), 'predResize': True, 'numpy': False, 'RGB_convert_first': True, 'mixFunction': { - 'function': mixSpillage_postprocess, + 'function': mixTraffic_postprocess, 'pars': { - #'modelSize': (640, 360), - 'modelSize': (1920,1080), + 'modelSize': (640, 360), 'RoadArea': 16000, 'roadVehicleAngle': 15, 'speedRoadVehicleAngleMax': 75, 'roundness': 1.0, - 'cls': 1, + 'cls': 0, 'vehicleFactor': 0.1, 'confThres': 0.25, 'roadIou': 0.6, @@ -831,26 +817,24 @@ class ModelType(Enum): 'labelnames': ["危化品","罐体","危险标识","普通车"], 'trtFlag_seg': True, 'trtFlag_det': True, - 'seg_nclass': 2, + 'seg_nclass': 3, 'segRegionCnt': 2, 'segPar': { - #'modelSize': (640, 360), - 'modelSize': (1920, 1080), + 'modelSize': (640, 360), 'mean': (0.485, 0.456, 0.406), 'std': (0.229, 0.224, 0.225), 'predResize': True, 'numpy': False, 'RGB_convert_first': True, 'mixFunction': { - 'function': mixCthc_postprocess, + 'function': mixTraffic_postprocess, 'pars': { - #'modelSize': (640, 360), - 'modelSize': (1920,1080), + 'modelSize': (640, 360), 'RoadArea': 16000, 'roadVehicleAngle': 15, 'speedRoadVehicleAngleMax': 75, 'roundness': 1.0, - 'cls': 4, + 'cls': 0, 'vehicleFactor': 0.1, 'confThres': 0.25, 'roadIou': 0.6, @@ -864,7 +848,7 @@ class ModelType(Enum): "name": "post_process", "conf_thres": 0.25, "iou_thres": 0.25, - "classes": 1, + "classes": 4, "rainbows": COLOR }, 'detModelpara': [{"id": str(x), "config": {"k1": "v1", "k2": "v2"}} for x in [0]], @@ -895,34 +879,44 @@ class ModelType(Enum): }) CITY_CARPLATE_MODEL = ("30", "301", "自研车牌检测", 'carplate', lambda device, gpuName: { + 'labelnames': ["车牌"], 'device': str(device), - 'models':{ + 'rainbows': COLOR, + 'models': [ { - 'weights': '../AIlib2/weights/conf/jkm/plate_yolov5s_v3.jit', - 'conf_thres': 0.4, - 'iou_thres': 0.45, - 'nc':1, + 'trtFlag_det': False, + 'weight': '../weights/pth/AIlib2/carplate/plate_yolov5s_v3.jit', + 'name': 'yolov5', + 'model': yolov5Model, + 'par': { + 'device': 'cuda:0', + 'half': False, + 'conf_thres': 0.4, + 'iou_thres': 0.45, + 'nc': 1, + 'plate_dilate': (0.5, 0.1) + }, }, { - 'weight' : '../weights/pth/AIlib2/ocr2/crnn_ch.pth', - 'name':'ocr', - 'model':ocrModel, - 'par':{ - 'char_file':'../AIlib2/conf/ocr2/benchmark.txt', - 'mode':'ch', - 'nc':3, - 'imgH':32, - 'imgW':192, - 'hidden':256, - 'mean':[0.5,0.5,0.5], - 'std':[0.5,0.5,0.5], - 'dynamic':False, - } - }, - } - }) + 'trtFlag_ocr': False, + 'weight': '../weights/pth/AIlib2/ocr2/crnn_ch.pth', + 'name': 'ocr', + 'model': ocrModel, + 'par': { + 'char_file': '../AIlib2/conf/ocr2/benchmark.txt', + 'mode': 'ch', + 'nc': 3, + 'imgH': 32, + 'imgW': 192, + 'hidden': 256, + 'mean': [0.5, 0.5, 0.5], + 'std': [0.5, 0.5, 0.5], + 'dynamic': False, + } + }], + }) - CITY_INFRAREDPERSON_MODEL = ("30", "302", "红外行人模型", 'infraredperson', lambda device, gpuName: { + CITY_INFRAREDPERSON_MODEL = ("30", "302", "红外行人模型", 'infraredPerson', lambda device, gpuName: { 'labelnames': ["行人"], 'postProcess': {'function': default_mix, 'pars': {}}, 'models': @@ -962,7 +956,7 @@ class ModelType(Enum): "rainbows": COLOR }, - }) + }) @staticmethod def checkCode(code):