更新 enums/ModelTypeEnum.py
This commit is contained in:
parent
9618bbc526
commit
919d15ec5f
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Reference in New Issue