更新 enums/ModelTypeEnum.py

This commit is contained in:
zhoushuliang 2025-07-05 11:13:00 +08:00
parent 7a2616df6a
commit 9906a10a66
1 changed files with 53 additions and 59 deletions

View File

@ -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):