From 919d15ec5ff55407c678ffdec134b7fecf362a47 Mon Sep 17 00:00:00 2001 From: zhoushuliang Date: Thu, 10 Jul 2025 17:21:15 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=20enums/ModelTypeEnum.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- enums/ModelTypeEnum.py | 112 +++++++++++++++++++++++++++++------------ 1 file changed, 80 insertions(+), 32 deletions(-) diff --git a/enums/ModelTypeEnum.py b/enums/ModelTypeEnum.py index 1f89471..fada509 100644 --- a/enums/ModelTypeEnum.py +++ b/enums/ModelTypeEnum.py @@ -16,6 +16,7 @@ from utilsK.illParkingUtils import illParking_postprocess from utilsK.pannelpostUtils import pannel_post_process from stdc import stdcModel from yolov5 import yolov5Model +from p2pNet import p2NnetModel from DMPRUtils.jointUtil import dmpr_yolo_stdc from AI import default_mix from ocr import ocrModel @@ -66,7 +67,7 @@ class ModelType(Enum): 'Detweights': "../weights/trt/AIlib2/river/yolov5_%s_fp16.engine" % gpuName, 'Segweights': '../weights/trt/AIlib2/river/stdc_360X640_%s_fp16.engine' % gpuName }) - + # FOREST_FARM_MODEL = ("2", "002", "森林模型", 'forest2', lambda device, gpuName: { # 'device': device, # 'gpu_name': gpuName, @@ -87,7 +88,7 @@ class ModelType(Enum): # }, # 'Segweights': None # }) - + FOREST_FARM_MODEL = ("2", "002", "森林模型", 'forest2', lambda device, gpuName: { 'labelnames': ["林斑", "病死树", "行人", "火焰", "烟雾","云朵"], @@ -100,10 +101,10 @@ class ModelType(Enum): '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.3,"2":0.3,"3":0.3 } }, } - - + + ], - + 'postFile': { "name": "post_process", "conf_thres": 0.25, @@ -115,7 +116,7 @@ class ModelType(Enum): 'segRegionCnt':2,###分割模型结果需要保留的等值线数目 "pixScale": 1.2, - + }) @@ -361,14 +362,14 @@ class ModelType(Enum): }) CITY_MANGEMENT_MODEL = ("16", "016", "城管模型", 'cityMangement2', lambda device, gpuName: { - 'labelnames': [ "车辆", "垃圾", "商贩", "违停","占道经营","裸土","未覆盖裸土" ], + 'labelnames': [ "车辆", "垃圾", "商贩", "违停","占道经营","裸土","未覆盖裸土","违建" ], 'postProcess':{ 'function':dmpr_yolo_stdc, 'pars':{ - 'carCls':0 ,'illCls':6,'scaleRatio':0.5,'border':80, - #车辆","垃圾","商贩","裸土","占道经营","违停"---> - #"车辆","垃圾","商贩","违停","占道经营","裸土" - 'classReindex':{ 0:0,1:1,2:2,3:6,4:4,5:5,6:3} + 'carCls':0 ,'illCls':7,'scaleRatio':0.5,'border':80, + #"车辆","垃圾","商贩","裸土","占道经营","未覆盖裸土","违建" + # key:实际训练index value:展示index + 'classReindex':{ 0:0,1:1,2:2,7:3,4:4,3:5,5:6,6:7} } }, 'models':[ @@ -376,7 +377,7 @@ class ModelType(Enum): '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':True,'trtFlag_seg':True, "score_byClass":{"0":0.8,"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,4,5,6,7],'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', @@ -384,25 +385,25 @@ class ModelType(Enum): 'depth_factor':32,'NUM_FEATURE_MAP_CHANNEL':6,'dmpr_thresh':0.1, 'dmprimg_size':640, 'name':'dmpr' }, - 'model':DMPRModel, + 'model':DMPRModel, 'name':'dmpr' }, - { + { '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':3},###分割模型预处理参数 'model':stdcModel, - 'name':'stdc' - } + 'name':'stdc' + } ], 'postFile': { "name": "post_process", "conf_thres": 0.25, "iou_thres": 0.45, - "classes": 6, + "classes": 8, "rainbows": COLOR }, - 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9] ],###控制哪些检测类别显示、输出 + 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6,7,8,9] ],###控制哪些检测类别显示、输出 'segRegionCnt':2,###分割模型结果需要保留的等值线数目 "pixScale": 1.2, }) @@ -569,7 +570,7 @@ class ModelType(Enum): '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/pth/AIlib2/ocr2/crnn_ch.pth', 'name':'ocr', 'model':ocrModel, @@ -584,7 +585,7 @@ class ModelType(Enum): 'std':[0.5,0.5,0.5], 'dynamic':False, }, - } + } ], 'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6]], 'segPar': None, @@ -597,7 +598,7 @@ class ModelType(Enum): }, 'Segweights': None, }) - + RIVERT_MODEL = ("25", "025", "河道检测模型(T)", 'riverT', lambda device, gpuName: { 'device': device, 'labelnames': ["漂浮物", "岸坡垃圾", "排口", "违建", "菜地", "水生植物", "河湖人员", "钓鱼人员", "船只", @@ -664,7 +665,8 @@ class ModelType(Enum): }) TRAFFICFORDSJ_FARM_MODEL = ("27", "027", "交通模型-大数据局", 'highWay2T', lambda device, gpuName: { 'device': str(device), - 'labelnames': ["行人", "车辆", "纵向裂缝", "横向裂缝", "修补", "网状裂纹", "坑槽", "块状裂纹", "积水", "影子", "事故", "桥梁外观","设施破损缺失","龙门架","防抛网","标识牌损坏","护栏损坏","钢筋裸露" ], + 'labelnames': ["行人", "车辆", "纵向裂缝", "横向裂缝", "修补", "网状裂纹", "坑槽", "块状裂纹", "积水", "影子", + "事故", "桥梁外观","设施破损缺失","龙门架","防抛网","标识牌损坏","护栏损坏","钢筋裸露" ], 'trtFlag_seg': True, 'trtFlag_det': True, 'seg_nclass': 3, @@ -718,13 +720,13 @@ class ModelType(Enum): 'model':yolov5Model, 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':list(range(20)),'segRegionCnt':1, 'trtFlag_det':True,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } }, } - - + + ], 'postFile': { "rainbows": COLOR }, - + }) RUBBISH_MODEL = ("29", "029", "垃圾模型", 'rubbish', lambda device, gpuName: { @@ -738,15 +740,15 @@ class ModelType(Enum): 'model':yolov5Model, 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':list(range(20)),'segRegionCnt':1, 'trtFlag_det':True,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } }, } - - + + ], 'postFile': { "rainbows": COLOR }, - + }) - + FIREWORK_MODEL = ("30", "030", "烟花模型", 'firework', lambda device, gpuName: { 'labelnames': [ "烟花"], 'postProcess':{'function':default_mix,'pars':{}}, @@ -758,13 +760,13 @@ class ModelType(Enum): 'model':yolov5Model, 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':list(range(20)),'segRegionCnt':1, 'trtFlag_det':True,'trtFlag_seg':False, "score_byClass":{"0":0.25,"1":0.3,"2":0.3,"3":0.3 } }, } - - + + ], 'postFile': { "rainbows": COLOR }, - + }) TRAFFIC_SPILL_MODEL = ("50", "501", "高速公路抛洒物模型", 'highWaySpill', lambda device, gpuName: { @@ -958,6 +960,52 @@ class ModelType(Enum): }) + CITY_DENSECROWDCOUNT_MODEL = ("30", "304", "密集人群计数", 'DenseCrowdCount', lambda device, gpuName: { + 'labelnames': ["人群计数"], + 'device': str(device), + 'rainbows': COLOR, + 'models': [ + { + 'trtFlag_det': False, + 'weight': "../weights/pth/AIlib2/DenseCrowd/SHTechA.pth", ###检测模型路径 + 'vggweight': "../weights/pth/AIlib2/DenseCrowd/vgg16_bn-6c64b313.pth", ###检测模型路径 + 'name': 'p2pnet', + 'model': p2NnetModel, + 'par': { + 'device': 'cuda:0', + 'row': 2, + 'line': 2, + 'point_loss_coef': 0.45, + 'conf': 0.25, + 'gpu_id': 0, + 'eos_coef': '0.5', + 'set_cost_class': 1, + 'set_cost_point': 0.05, + 'backbone': 'vgg16_bn' + }, + }], + }) + + CITY_DENSECROWDESTIMATION_MODEL = ("30", "305", "密集人群密度估计", 'DenseCrowdEstimation', lambda device, gpuName: { + 'labelnames': ["密度"], + 'models': + [ + { + 'weight': "../weights/pth/AIlib2/DenseCrowd/SHTechA.pth", ###检测模型路径 + 'name': 'yolov5', + 'model': yolov5Model, + 'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.50, 'iou_thres': 0.45, + 'allowedList': list(range(20)), 'segRegionCnt': 1, 'trtFlag_det': True, + 'trtFlag_seg': False, "score_byClass": {"0": 0.50, "1": 0.3, "2": 0.3, "3": 0.3}}, + } + + ], + 'postFile': { + "rainbows": COLOR + }, + + }) + @staticmethod def checkCode(code): for model in ModelType: