gitadmin
/
tuoheng_algN
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

更新 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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

70
enums/ModelTypeEnum.py
View File

@ -9,12 +9,10 @@ from DMPRUtils.jointUtil import dmpr_yolo
from segutils.segmodel import SegModel from segutils.segmodel import SegModel
from utilsK.queRiver import riverDetSegMixProcess from utilsK.queRiver import riverDetSegMixProcess
from utilsK.crowdGather import gather_post_process 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.drownUtils import mixDrowing_water_postprocess
from utilsK.noParkingUtils import mixNoParking_road_postprocess from utilsK.noParkingUtils import mixNoParking_road_postprocess
from utilsK.illParkingUtils import illParking_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 utilsK.pannelpostUtils import pannel_post_process
from stdc import stdcModel from stdc import stdcModel
from yolov5 import yolov5Model from yolov5 import yolov5Model
@ -121,7 +119,6 @@ class ModelType(Enum):
}) })
TRAFFIC_FARM_MODEL = ("3", "003", "交通模型", 'highWay2', lambda device, gpuName: { TRAFFIC_FARM_MODEL = ("3", "003", "交通模型", 'highWay2', lambda device, gpuName: {
'device': str(device), 'device': str(device),
'labelnames': ["行人", "车辆", "纵向裂缝", "横向裂缝", "修补", "网状裂纹", "坑槽", "块状裂纹", "积水", "影子", 'labelnames': ["行人", "车辆", "纵向裂缝", "横向裂缝", "修补", "网状裂纹", "坑槽", "块状裂纹", "积水", "影子",
@ -359,9 +356,7 @@ class ModelType(Enum):
"classes": 5, "classes": 5,
"rainbows": COLOR "rainbows": COLOR
}, },
# "../weights/pth/AIlib2/%s/yolov5.pt" % modeType.value[3]
'Detweights': "../weights/trt/AIlib2/river2/yolov5_%s_fp16.engine" % gpuName, '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 'Segweights': '../weights/trt/AIlib2/river2/stdc_360X640_%s_fp16.engine' % gpuName
}) })
@ -378,11 +373,10 @@ class ModelType(Enum):
}, },
'models':[ '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', 'name':'yolov5',
'model':yolov5Model, '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', 'weight':'../weights/pth/AIlib2/cityMangement3/dmpr.pth',
@ -394,9 +388,9 @@ class ModelType(Enum):
'name':'dmpr' 'name':'dmpr'
}, },
{ {
'weight':'../weights/pth/AIlib2/cityMangement3/stdc_360X640.pth', 'weight':'../weights/trt/AIlib2/cityMangement3/stdc_360X640_%s_fp16.engine'%(gpuName),
'par':{ '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, 'model':stdcModel,
'name':'stdc' 'name':'stdc'
} }
@ -441,9 +435,7 @@ class ModelType(Enum):
"classes": 9, "classes": 9,
"rainbows": COLOR "rainbows": COLOR
}, },
# "../weights/pth/AIlib2/%s/yolov5.pt" % modeType.value[3]
'Detweights': "../weights/trt/AIlib2/drowning/yolov5_%s_fp16.engine" % gpuName, '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 '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: { CITYROAD_MODEL = ("20", "020", "城市公路模型", 'cityRoad', lambda device, gpuName: {
'device': device, 'device': device,
'labelnames': ["护栏", "交通标志", "非交通标志", "施工", "施工"], 'labelnames': ["护栏", "交通标志", "非交通标志", "施工锥桶", "施工水马"],
'trtFlag_seg': False, 'trtFlag_seg': False,
'trtFlag_det': True, 'trtFlag_det': True,
'slopeIndex': [], 'slopeIndex': [],
@ -572,14 +564,12 @@ class ModelType(Enum):
}}, }},
'models':[ 'models':[
{ {
#'weight':'../weights/pth/AIlib2/channel2/yolov5.pt',
'weight':'../weights/trt/AIlib2/channel2/yolov5_%s_fp16.engine'%(gpuName), 'weight':'../weights/trt/AIlib2/channel2/yolov5_%s_fp16.engine'%(gpuName),
'name':'yolov5', 'name':'yolov5',
'model':yolov5Model, '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} } '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', 'weight' : '../weights/pth/AIlib2/ocr2/crnn_ch.pth',
'name':'ocr', 'name':'ocr',
'model':ocrModel, 'model':ocrModel,
@ -638,9 +628,7 @@ class ModelType(Enum):
"classes": 5, "classes": 5,
"rainbows": COLOR "rainbows": COLOR
}, },
# "../weights/pth/AIlib2/%s/yolov5.pt" % modeType.value[3]
'Detweights': "../weights/trt/AIlib2/riverT/yolov5_%s_fp16.engine" % gpuName, '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 'Segweights': '../weights/trt/AIlib2/riverT/stdc_360X640_%s_fp16.engine' % gpuName
}) })
@ -784,26 +772,24 @@ class ModelType(Enum):
'labelnames': ["抛洒物","车辆"], 'labelnames': ["抛洒物","车辆"],
'trtFlag_seg': True, 'trtFlag_seg': True,
'trtFlag_det': True, 'trtFlag_det': True,
'seg_nclass': 2, 'seg_nclass': 3,
'segRegionCnt': 2, 'segRegionCnt': 2,
'segPar': { 'segPar': {
#'modelSize': (640, 360), 'modelSize': (640, 360),
'modelSize': (1920, 1080),
'mean': (0.485, 0.456, 0.406), 'mean': (0.485, 0.456, 0.406),
'std': (0.229, 0.224, 0.225), 'std': (0.229, 0.224, 0.225),
'predResize': True, 'predResize': True,
'numpy': False, 'numpy': False,
'RGB_convert_first': True, 'RGB_convert_first': True,
'mixFunction': { 'mixFunction': {
'function': mixSpillage_postprocess, 'function': mixTraffic_postprocess,
'pars': { 'pars': {
#'modelSize': (640, 360), 'modelSize': (640, 360),
'modelSize': (1920,1080),
'RoadArea': 16000, 'RoadArea': 16000,
'roadVehicleAngle': 15, 'roadVehicleAngle': 15,
'speedRoadVehicleAngleMax': 75, 'speedRoadVehicleAngleMax': 75,
'roundness': 1.0, 'roundness': 1.0,
'cls': 1, 'cls': 0,
'vehicleFactor': 0.1, 'vehicleFactor': 0.1,
'confThres': 0.25, 'confThres': 0.25,
'roadIou': 0.6, 'roadIou': 0.6,
@ -831,26 +817,24 @@ class ModelType(Enum):
'labelnames': ["危化品","罐体","危险标识","普通车"], 'labelnames': ["危化品","罐体","危险标识","普通车"],
'trtFlag_seg': True, 'trtFlag_seg': True,
'trtFlag_det': True, 'trtFlag_det': True,
'seg_nclass': 2, 'seg_nclass': 3,
'segRegionCnt': 2, 'segRegionCnt': 2,
'segPar': { 'segPar': {
#'modelSize': (640, 360), 'modelSize': (640, 360),
'modelSize': (1920, 1080),
'mean': (0.485, 0.456, 0.406), 'mean': (0.485, 0.456, 0.406),
'std': (0.229, 0.224, 0.225), 'std': (0.229, 0.224, 0.225),
'predResize': True, 'predResize': True,
'numpy': False, 'numpy': False,
'RGB_convert_first': True, 'RGB_convert_first': True,
'mixFunction': { 'mixFunction': {
'function': mixCthc_postprocess, 'function': mixTraffic_postprocess,
'pars': { 'pars': {
#'modelSize': (640, 360), 'modelSize': (640, 360),
'modelSize': (1920,1080),
'RoadArea': 16000, 'RoadArea': 16000,
'roadVehicleAngle': 15, 'roadVehicleAngle': 15,
'speedRoadVehicleAngleMax': 75, 'speedRoadVehicleAngleMax': 75,
'roundness': 1.0, 'roundness': 1.0,
'cls': 4, 'cls': 0,
'vehicleFactor': 0.1, 'vehicleFactor': 0.1,
'confThres': 0.25, 'confThres': 0.25,
'roadIou': 0.6, 'roadIou': 0.6,
@ -864,7 +848,7 @@ class ModelType(Enum):
"name": "post_process", "name": "post_process",
"conf_thres": 0.25, "conf_thres": 0.25,
"iou_thres": 0.25, "iou_thres": 0.25,
"classes": 1, "classes": 4,
"rainbows": COLOR "rainbows": COLOR
}, },
'detModelpara': [{"id": str(x), "config": {"k1": "v1", "k2": "v2"}} for x in [0]], 'detModelpara': [{"id": str(x), "config": {"k1": "v1", "k2": "v2"}} for x in [0]],
@ -895,15 +879,26 @@ class ModelType(Enum):
}) })
CITY_CARPLATE_MODEL = ("30", "301", "自研车牌检测", 'carplate', lambda device, gpuName: { CITY_CARPLATE_MODEL = ("30", "301", "自研车牌检测", 'carplate', lambda device, gpuName: {
'labelnames': ["车牌"],
'device': str(device), 'device': str(device),
'models':{ 'rainbows': COLOR,
'models': [
{ {
'weights': '../AIlib2/weights/conf/jkm/plate_yolov5s_v3.jit', '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, 'conf_thres': 0.4,
'iou_thres': 0.45, 'iou_thres': 0.45,
'nc': 1, 'nc': 1,
'plate_dilate': (0.5, 0.1)
},
}, },
{ {
'trtFlag_ocr': False,
'weight': '../weights/pth/AIlib2/ocr2/crnn_ch.pth', 'weight': '../weights/pth/AIlib2/ocr2/crnn_ch.pth',
'name': 'ocr', 'name': 'ocr',
'model': ocrModel, 'model': ocrModel,
@ -918,11 +913,10 @@ class ModelType(Enum):
'std': [0.5, 0.5, 0.5], 'std': [0.5, 0.5, 0.5],
'dynamic': False, 'dynamic': False,
} }
}, }],
}
}) })
CITY_INFRAREDPERSON_MODEL = ("30", "302", "红外行人模型", 'infraredperson', lambda device, gpuName: { CITY_INFRAREDPERSON_MODEL = ("30", "302", "红外行人模型", 'infraredPerson', lambda device, gpuName: {
'labelnames': ["行人"], 'labelnames': ["行人"],
'postProcess': {'function': default_mix, 'pars': {}}, 'postProcess': {'function': default_mix, 'pars': {}},
'models': 'models':