tuoheng_algN/enums/ModelTypeEnum.py

1138 lines
41 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import sys
from enum import Enum, unique
from common.Constant import COLOR
sys.path.extend(['..', '../AIlib2'])
from DMPR import DMPRModel
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,mixTraffic_postprocess
from utilsK.drownUtils import mixDrowing_water_postprocess
from utilsK.noParkingUtils import mixNoParking_road_postprocess
from utilsK.illParkingUtils import illParking_postprocess
from utilsK.pannelpostUtils import pannel_post_process
from utilsK.securitypostUtils import security_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
from utilsK.channel2postUtils import channel2_post_process
'''
参数说明
1. 编号
2. 模型编号
3. 模型名称
4. 选用的模型名称
5. 模型配置
6. 模型引用配置[Detweights文件, Segweights文件, 引用计数]
'''
@unique
class ModelType(Enum):
WATER_SURFACE_MODEL = ("1", "001", "河道模型", 'river', lambda device, gpuName: {
'device': device,
'labelnames': ["排口", "水生植被", "其它", "漂浮物", "污染排口", "菜地", "违建", "岸坡垃圾"],
'seg_nclass': 2,
'trtFlag_seg': True,
'trtFlag_det': True,
'segRegionCnt': 1,
'segPar': {
'modelSize': (640, 360),
'mean': (0.485, 0.456, 0.406),
'std': (0.229, 0.224, 0.225),
'numpy': False,
'RGB_convert_first': True,
'mixFunction': {
'function': riverDetSegMixProcess,
'pars': {
'slopeIndex': [5, 6, 7],
'riverIou': 0.1
}
}
},
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
'fiterList':[2],
'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,
# 'labelnames': ["林斑", "病死树", "行人", "火焰", "烟雾","云朵"],
# 'trtFlag_det': True,
# 'trtFlag_seg': False,
# 'Detweights': "../weights/trt/AIlib2/forest2/yolov5_%s_fp16.engine" % gpuName,
# 'seg_nclass': 2,
# 'segRegionCnt': 0,
# 'slopeIndex': [],
# 'segPar': None,
# 'postFile': {
# "name": "post_process",
# "conf_thres": 0.25,
# "iou_thres": 0.45,
# "classes": 6,
# "rainbows": COLOR
# },
# 'Segweights': None
# })
FOREST_FARM_MODEL = ("2", "002", "森林模型", 'forest2', lambda device, gpuName: {
'labelnames': ["林斑", "病死树", "行人", "火焰", "烟雾","云朵"],
'postProcess':{'function':default_mix,'pars':{}},
'models':
[
{
'weight':"../weights/trt/AIlib2/forest2/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径
'name':'yolov5',
'model':yolov5Model,
'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False},
}
],
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
"score_byClass": {0: 0.25, 1: 0.3, 2: 0.3, 3: 0.3},
'fiterList': [5],
'segRegionCnt':2,###分割模型结果需要保留的等值线数目
"pixScale": 1.2,
})
TRAFFIC_FARM_MODEL = ("3", "003", "交通模型", 'highWay2', lambda device, gpuName: {
'device': str(device),
'labelnames': ["行人", "车辆", "纵向裂缝", "横向裂缝", "修补", "网状裂纹", "坑槽", "块状裂纹", "积水", "影子",
"事故","抛撒物", "危化品车辆", "虚标线","其他标线","其他","桥梁外观","设施破损缺失","龙门架","防抛网","标识牌损坏","护栏损坏","钢筋裸露"],
'trtFlag_seg': True,
'trtFlag_det': True,
'seg_nclass': 3,
'segRegionCnt': 2,
'segPar': {
'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': tracfficAccidentMixFunction,
'pars': {
'modelSize': (640, 360),
'RoadArea': 16000,
'roadVehicleAngle': 15,
'speedRoadVehicleAngleMax': 75,
'roundness': 1.0,
'cls': 10,
'CarId':1,
'CthcId':12,
'vehicleFactor': 0.1,
'confThres': 0.25,
'roadIou': 0.6,
'radius': 50,
'vehicleFlag': False,
'distanceFlag': False
}
}
},
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.25,
"classes": 10,
"rainbows": COLOR
},
'score_byClass':{11:0.75,12:0.75},
'fiterList': [13,14,15,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
})
EPIDEMIC_PREVENTION_MODEL = ("4", "004", "防疫模型", None, None)
PLATE_MODEL = ("5", "005", "车牌模型", None, None)
VEHICLE_MODEL = ("6", "006", "车辆模型", 'vehicle', lambda device, gpuName: {
'device': device,
'gpu_name': gpuName,
'labelnames': ["车辆"],
'seg_nclass': 2,
'segRegionCnt': 0,
'slopeIndex': [],
'trtFlag_det': True,
'trtFlag_seg': False,
'Detweights': "../weights/trt/AIlib2/vehicle/yolov5_%s_fp16.engine" % gpuName,
'segPar': None,
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
'Segweights': None
})
PEDESTRIAN_MODEL = ("7", "007", "行人模型", 'pedestrian', lambda device, gpuName: {
'device': device,
'gpu_name': gpuName,
'labelnames': ["行人"],
'seg_nclass': 2,
'segRegionCnt': 0,
'trtFlag_det': True,
'trtFlag_seg': False,
'Detweights': "../weights/trt/AIlib2/pedestrian/yolov5_%s_fp16.engine" % gpuName,
'slopeIndex': [],
'segPar': None,
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
'Segweights': None
})
SMOGFIRE_MODEL = ("8", "008", "烟火模型", 'smogfire', lambda device, gpuName: {
'device': device,
'gpu_name': gpuName,
'labelnames': ["火焰", "烟雾"],
'seg_nclass': 2, # 分割模型类别数目默认2类
'segRegionCnt': 0,
'trtFlag_det': True,
'trtFlag_seg': False,
'Detweights': "../weights/trt/AIlib2/smogfire/yolov5_%s_fp16.engine" % gpuName,
'slopeIndex': [],
'segPar': None,
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
'Segweights': None,
})
ANGLERSWIMMER_MODEL = ("9", "009", "钓鱼游泳模型", 'AnglerSwimmer', lambda device, gpuName: {
'device': device,
'gpu_name': gpuName,
'labelnames': ["钓鱼", "游泳"],
'seg_nclass': 2, # 分割模型类别数目默认2类
'segRegionCnt': 0,
'slopeIndex': [],
'trtFlag_det': True,
'trtFlag_seg': False,
'Detweights': "../weights/trt/AIlib2/AnglerSwimmer/yolov5_%s_fp16.engine" % gpuName,
'segPar': None,
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
'Segweights': None
})
COUNTRYROAD_MODEL = ("10", "010", "乡村模型", 'countryRoad', lambda device, gpuName: {
'device': device,
'gpu_name': gpuName,
'labelnames': ["违法种植"],
'seg_nclass': 2, # 分割模型类别数目默认2类
'segRegionCnt': 0,
'slopeIndex': [],
'trtFlag_det': True,
'trtFlag_seg': False,
'Detweights': "../weights/trt/AIlib2/countryRoad/yolov5_%s_fp16.engine" % gpuName,
'segPar': None,
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
'Segweights': None
})
SHIP_MODEL = ("11", "011", "船只模型", 'ship2', lambda device, gpuName: {
'model_size': (608, 608),
'K': 100,
'conf_thresh': 0.18,
'device': 'cuda:%s' % device,
'down_ratio': 4,
'num_classes': 15,
'weights': '../weights/trt/AIlib2/ship2/obb_608X608_%s_fp16.engine' % gpuName,
'dataset': 'dota',
'half': False,
'mean': (0.5, 0.5, 0.5),
'std': (1, 1, 1),
'heads': {'hm': None, 'wh': 10, 'reg': 2, 'cls_theta': 1},
'decoder': None,
'test_flag': True,
"rainbows": COLOR,
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
'drawBox': False,
'label_array': None,
'labelnames': ("0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "船只"),
})
BAIDU_MODEL = ("12", "012", "百度AI图片识别模型", None, None)
CHANNEL_EMERGENCY_MODEL = ("13", "013", "航道模型", 'channelEmergency', lambda device, gpuName: {
'device': device,
'gpu_name': gpuName,
'labelnames': [""],
'seg_nclass': 2, # 分割模型类别数目默认2类
'segRegionCnt': 0,
'slopeIndex': [],
'trtFlag_det': True,
'trtFlag_seg': False,
'Detweights': "../weights/trt/AIlib2/channelEmergency/yolov5_%s_fp16.engine" % gpuName,
'segPar': None,
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
'Segweights': None
})
RIVER2_MODEL = ("15", "015", "河道检测模型", 'river2', lambda device, gpuName: {
'device': device,
'labelnames': ["漂浮物", "岸坡垃圾", "排口", "违建", "菜地", "水生植物", "河湖人员", "钓鱼人员", "船只",
"蓝藻"],
'trtFlag_seg': True,
'trtFlag_det': True,
'seg_nclass': 2,
'segRegionCnt': 1,
'segPar': {
'modelSize': (640, 360),
'mean': (0.485, 0.456, 0.406),
'std': (0.229, 0.224, 0.225),
'numpy': False,
'RGB_convert_first': True,
'mixFunction': {
'function': riverDetSegMixProcess,
'pars': {
'slopeIndex': [1, 3, 4, 7],
'riverIou': 0.1,
'scale': 0.25
}
}
},
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.3,
"ovlap_thres_crossCategory": 0.65,
"classes": 5,
"rainbows": COLOR
},
'Detweights': "../weights/trt/AIlib2/river2/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/river2/stdc_360X640_%s_fp16.engine' % gpuName
})
CITY_MANGEMENT_MODEL = ("16", "016", "城管模型", 'cityMangement2', lambda device, gpuName: {
'labelnames': [ "车辆", "垃圾", "商贩", "违停","占道经营","裸土","未覆盖裸土","违建" ],
'postProcess':{
'function':dmpr_yolo_stdc,
'pars':{
'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':[
{
'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,'segRegionCnt':1, 'trtFlag_det':True,'trtFlag_seg':True}
},
{
'weight':'../weights/trt/AIlib2/cityMangement3/dmpr_3090.engine',
#'weight':'../weights/pth/AIlib2/cityMangement3/dmpr.pth',
'par':{
'depth_factor':32,'NUM_FEATURE_MAP_CHANNEL':6,'dmpr_thresh':0.1, 'dmprimg_size':640,
'name':'dmpr'
},
'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'
}
],
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 8,
"rainbows": COLOR
},
"score_byClass":{0:0.8, 1:0.4, 2:0.5, 3:0.5},
'segRegionCnt':2,###分割模型结果需要保留的等值线数目
"pixScale": 1.2,
})
DROWING_MODEL = ("17", "017", "人员落水模型", 'drowning', lambda device, gpuName: {
'device': device,
'labelnames': ["人头", "", "船只"],
'trtFlag_seg': True,
'trtFlag_det': True,
'seg_nclass': 2,
'segRegionCnt': 2,
'segPar': {
'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': mixDrowing_water_postprocess,
'pars': {
'modelSize': (640, 360)
}
}
},
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.25,
"classes": 9,
"rainbows": COLOR
},
'Detweights': "../weights/trt/AIlib2/drowning/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/drowning/stdc_360X640_%s_fp16.engine' % gpuName
})
NOPARKING_MODEL = (
"18", "018", "城市违章模型", 'noParking', lambda device, gpuName: {
'device': device,
'labelnames': ["车辆", "违停"],
'trtFlag_seg': True,
'trtFlag_det': True,
'seg_nclass': 4,
'segRegionCnt': 2,
'segPar': {
'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': mixNoParking_road_postprocess,
'pars': {
'modelSize': (640, 360),
'roundness': 0.3,
'cls': 9,
'laneArea': 10,
'laneAngleCha': 5,
'RoadArea': 16000,
'fitOrder':2
}
}
},
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.25,
"classes": 9,
"rainbows": COLOR
},
'Detweights': "../weights/trt/AIlib2/noParking/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/noParking/stdc_360X640_%s_fp16.engine' % gpuName
})
ILLPARKING_MODEL = ("19", "019", "车辆违停模型", 'illParking', lambda device, gpuName: {
'device': device,
'labelnames': ["", "T角点", "L角点", "违停"],
'trtFlag_seg': False,
'trtFlag_det': True,
'seg_nclass': 4,
'segRegionCnt': 2,
'segPar': {
'mixFunction': {
'function': illParking_postprocess,
'pars': {}
}
},
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.25,
"classes": 9,
"rainbows": COLOR
},
'Detweights': "../weights/trt/AIlib2/illParking/yolov5_%s_fp16.engine" % gpuName,
'Segweights': None
})
CITYROAD_MODEL = ("20", "020", "城市公路模型", 'cityRoad', lambda device, gpuName: {
'device': device,
'labelnames': ["护栏", "交通标志", "非交通标志", "施工锥桶", "施工水马"],
'trtFlag_seg': False,
'trtFlag_det': True,
'slopeIndex': [],
'seg_nclass': 2,
'segRegionCnt': 0,
'segPar': None,
'postFile': {
"name": "post_process",
"conf_thres": 0.8,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
'Detweights': "../weights/trt/AIlib2/cityRoad/yolov5_%s_fp16.engine" % gpuName,
'Segweights': None
})
POTHOLE_MODEL = ("23", "023", "坑槽检测模型", 'pothole', lambda device, gpuName: {
'device': device,
'gpu_name': gpuName,
'labelnames': ["坑槽"],
'seg_nclass': 2, # 分割模型类别数目默认2类
'segRegionCnt': 0,
'slopeIndex': [],
'trtFlag_det': True,
'trtFlag_seg': False,
'Detweights': "../weights/trt/AIlib2/pothole/yolov5_%s_fp16.engine" % gpuName,
'segPar': None,
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
'Segweights': None,
})
CHANNEL2_MODEL = ("24", "024", "船只综合检测模型", 'channel2', lambda device, gpuName: {
'device': device,
'gpu_name': gpuName,
# 'labelnames': ["国旗", "浮标", "船名", "船只","未挂国旗船只"],
'labelnames': ["国旗", "浮标", "船名", "船只", "未挂国旗船只","未封仓船只","未挂国旗且未封仓船只"],
'segRegionCnt': 0,
'postProcess':{'function':channel2_post_process,'name':'channel2','pars':{
'objs':[2],
'wRation':1/6.0,
'hRation':1/6.0,
'flagId':0,
'boatId':3,
'unflagId': 4, # 未挂国旗船只
'uncoverId': 5, # 未封仓
'unflagAndcoverId': 6, # 未挂国旗且未封仓
'recScale':1.2,
'target_cls': 3, # 船只目标种类
'filter_cls': 4 # 被过滤的种类,模型文件中未封仓实际index
}},
'models':[
{
'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,'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False}
},
{
'weight' : '../weights/trt/AIlib2/ocr2/crnn_ch_%s_fp16_192X32.engine'%(gpuName),
'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,
},
}
],
'segPar': None,
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
'Segweights': None,
"score_byClass": {0: 0.7, 1: 0.7, 2: 0.8, 3: 0.6}
})
RIVERT_MODEL = ("25", "025", "河道检测模型(T)", 'riverT', lambda device, gpuName: {
'device': device,
'labelnames': ["漂浮物", "岸坡垃圾", "排口", "违建", "菜地", "水生植物", "河湖人员", "钓鱼人员", "船只",
"蓝藻"],
'trtFlag_seg': True,
'trtFlag_det': True,
'seg_nclass': 2,
'segRegionCnt': 1,
'segPar': {
'modelSize': (640, 360),
'mean': (0.485, 0.456, 0.406),
'std': (0.229, 0.224, 0.225),
'numpy': False,
'RGB_convert_first': True,
'mixFunction': {
'function': riverDetSegMixProcess,
'pars': {
'slopeIndex': [1, 3, 4, 7],
'riverIou': 0.1
}
}
},
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.3,
"ovlap_thres_crossCategory": 0.65,
"classes": 5,
"rainbows": COLOR
},
'Detweights': "../weights/trt/AIlib2/riverT/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/riverT/stdc_360X640_%s_fp16.engine' % gpuName
})
FORESTCROWD_FARM_MODEL = ("26", "026", "森林人群模型", 'forestCrowd', lambda device, gpuName: {
'labelnames': ["林斑", "病死树", "行人", "火焰", "烟雾","人群"],
'postProcess':{'function':gather_post_process,'pars':{'pedestrianId':2,'crowdThreshold':4,'gatherId':5,'distancePersonScale':2.0}},
'models':
[
{
'weight':"../weights/trt/AIlib2/forestCrowd/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径
'name':'yolov5',
'model':yolov5Model,
'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False},
}
],
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
"score_byClass":{0:0.25,1:0.25,2:0.6,3:0.6,4:0.6 ,5:0.6},
'segRegionCnt':2,###分割模型结果需要保留的等值线数目
"pixScale": 1.2,
})
TRAFFICFORDSJ_FARM_MODEL = ("27", "027", "交通模型-大数据局", 'highWay2T', lambda device, gpuName: {
'device': str(device),
'labelnames': ["行人", "车辆", "纵向裂缝", "横向裂缝", "修补", "网状裂纹", "坑槽", "块状裂纹", "积水", "影子",
"事故", "桥梁外观","设施破损缺失","龙门架","防抛网","标识牌损坏","护栏损坏","钢筋裸露" ],
'trtFlag_seg': True,
'trtFlag_det': True,
'seg_nclass': 3,
'segRegionCnt': 2,
'segPar': {
'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': tracfficAccidentMixFunction,
'pars': {
'modelSize': (640, 360),
'RoadArea': 16000,
'roadVehicleAngle': 15,
'speedRoadVehicleAngleMax': 75,
'roundness': 1.0,
'cls': 10,
'CarId':1,
'CthcId':1,
'vehicleFactor': 0.1,
'confThres': 0.25,
'roadIou': 0.6,
'radius': 50,
'vehicleFlag': False,
'distanceFlag': False
}
}
},
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.25,
"classes": 10,
"rainbows": COLOR
},
'fiterltList': [11,12,13,14,15,16,17],
'Detweights': "../weights/trt/AIlib2/highWay2T/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/highWay2T/stdc_360X640_%s_fp16.engine' % gpuName
})
SMARTSITE_MODEL = ("28", "028", "智慧工地模型", 'smartSite', lambda device, gpuName: {
'labelnames': [ "工人","塔式起重机","悬臂","起重机","压路机","推土机","挖掘机","卡车","装载机","泵车","混凝土搅拌车","打桩","其他车辆" ],
'postProcess':{'function':default_mix,'pars':{}},
'models':
[
{
'weight':"../weights/trt/AIlib2/smartSite/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径
'name':'yolov5',
'model':yolov5Model,
'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'segRegionCnt':1, 'trtFlag_det':True,'trtFlag_seg':False},
}
],
'postFile': {
"rainbows": COLOR
},
"score_byClass": {0: 0.25, 1: 0.3, 2: 0.3, 3: 0.3}
})
RUBBISH_MODEL = ("29", "029", "垃圾模型", 'rubbish', lambda device, gpuName: {
'labelnames': [ "建筑垃圾","白色垃圾","其他垃圾"],
'postProcess':{'function':default_mix,'pars':{}},
'models':
[
{
'weight':"../weights/trt/AIlib2/rubbish/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径
'name':'yolov5',
'model':yolov5Model,
'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'segRegionCnt':1, 'trtFlag_det':True,'trtFlag_seg':False},
}
],
'postFile': {
"rainbows": COLOR
},
"score_byClass": {0: 0.25, 1: 0.3, 2: 0.3, 3: 0.3}
})
FIREWORK_MODEL = ("30", "030", "烟花模型", 'firework', lambda device, gpuName: {
'labelnames': [ "烟花"],
'postProcess':{'function':default_mix,'pars':{}},
'models':
[
{
'weight':"../weights/trt/AIlib2/firework/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径
'name':'yolov5',
'model':yolov5Model,
'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'segRegionCnt':1, 'trtFlag_det':True,'trtFlag_seg':False },
}
],
'postFile': {
"rainbows": COLOR
},
})
TRAFFIC_SPILL_MODEL = ("50", "501", "高速公路抛洒物模型", 'highWaySpill', lambda device, gpuName: {
'device': str(device),
'labelnames': ["抛洒物","车辆"],
'trtFlag_seg': True,
'trtFlag_det': True,
'seg_nclass': 3,
'segRegionCnt': 2,
'segPar': {
'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': mixTraffic_postprocess,
'pars': {
'modelSize': (640, 360),
'RoadArea': 16000,
'roadVehicleAngle': 15,
'speedRoadVehicleAngleMax': 75,
'roundness': 1.0,
'cls': 0,
'vehicleFactor': 0.1,
'confThres': 0.25,
'roadIou': 0.6,
'radius': 50,
'vehicleFlag': False,
'distanceFlag': False
}
}
},
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.25,
"classes": 2,
"rainbows": COLOR
},
'fiterList': [1],
###控制哪些检测类别显示、输出
'Detweights': "../weights/trt/AIlib2/highWaySpill/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/highWaySpill/stdc_360X640_%s_fp16.engine' % gpuName
})
TRAFFIC_CTHC_MODEL = ("50", "502", "高速公路危化品模型", 'highWayCthc', lambda device, gpuName: {
'device': str(device),
'labelnames': ["危化品","罐体","危险标识","普通车"],
'trtFlag_seg': True,
'trtFlag_det': True,
'seg_nclass': 3,
'segRegionCnt': 2,
'segPar': {
'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': mixTraffic_postprocess,
'pars': {
'modelSize': (640, 360),
'RoadArea': 16000,
'roadVehicleAngle': 15,
'speedRoadVehicleAngleMax': 75,
'roundness': 1.0,
'cls': 0,
'vehicleFactor': 0.1,
'confThres': 0.25,
'roadIou': 0.6,
'radius': 50,
'vehicleFlag': False,
'distanceFlag': False
}
}
},
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.25,
"classes": 4,
"rainbows": COLOR
},
'fiterList':[1,2,3],
###控制哪些检测类别显示、输出
'Detweights': "../weights/trt/AIlib2/highWayCthc/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/highWayCthc/stdc_360X640_%s_fp16.engine' % gpuName
})
TRAFFIC_PANNEL_MODEL = ("50", "503", "光伏板模型", 'pannel', lambda device, gpuName: {
'labelnames': ["光伏板","覆盖物","裂缝"],
'postProcess': {'function': pannel_post_process, 'pars': {'objs': [0]}},
'models':
[
{
'weight': "../weights/trt/AIlib2/pannel/yolov5_%s_fp16.engine" % (gpuName), ###检测模型路径
'name': 'yolov5',
'model': yolov5Model,
'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.25, 'iou_thres': 0.45,
'segRegionCnt': 1, 'trtFlag_det': True,
'trtFlag_seg': False},
}
],
'postFile': {
"rainbows": COLOR
},
'fiterList':[0]
})
CITY_CARPLATE_MODEL = ("30", "301", "自研车牌检测", 'carplate', lambda device, gpuName: {
'labelnames': ["车牌"],
'device': str(device),
'rainbows': COLOR,
'models': [
{
#'weight': '../weights/pth/AIlib2/carplate/plate_yolov5s_v3.jit',
'weight': '../weights/trt/AIlib2/carplate/yolov5_%s_fp16.engine' % (gpuName),
'name': 'yolov5',
'model': yolov5Model,
'par': {
'trtFlag_det': True,
'device': 'cuda:0',
'half': True,
'conf_thres': 0.4,
'iou_thres': 0.45,
'nc': 1,
'plate':8,
'plate_dilate': (0.5, 0.1)
},
},
{
'weight' : '../weights/trt/AIlib2/ocr2/crnn_ch_%s_fp16_192X32.engine'%(gpuName),
'name': 'ocr',
'model': ocrModel,
'par': {
'trtFlag_ocr': True,
'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: {
'labelnames': ["行人"],
'postProcess': {'function': default_mix, 'pars': {}},
'models':
[
{
'weight': "../weights/trt/AIlib2/infraredPerson/yolov5_%s_fp16.engine" % (gpuName), ###检测模型路径
'name': 'yolov5',
'model': yolov5Model,
'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.50, 'iou_thres': 0.45,
'segRegionCnt': 1, 'trtFlag_det': True,'trtFlag_seg': False},
}
],
'postFile': {
"rainbows": COLOR
},
})
CITY_NIGHTFIRESMOKE_MODEL = ("30", "303", "夜间烟火模型", 'nightFireSmoke', lambda device, gpuName: {
'labelnames': ["","烟雾"],
'postProcess': {'function': default_mix, 'pars': {}},
'models':
[
{
'weight': "../weights/trt/AIlib2/nightFireSmoke/yolov5_%s_fp16.engine" % (gpuName), ###检测模型路径
'name': 'yolov5',
'model': yolov5Model,
'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.50, 'iou_thres': 0.45,
'segRegionCnt': 1, 'trtFlag_det': True, 'trtFlag_seg': False},
}
],
'postFile': {
"rainbows": COLOR
},
})
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.65,
'gpu_id': 0,
'eos_coef': '0.5',
'set_cost_class': 1,
'set_cost_point': 0.05,
'backbone': 'vgg16_bn',
'expend': 10,
'psize': 2,
},
}],
})
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,
'segRegionCnt': 1, 'trtFlag_det': True, 'trtFlag_seg': False},
}
],
'postFile': {
"rainbows": COLOR
},
})
CITY_UNDERBUILDCOUNT_MODEL = ("30", "306", "建筑物下人群计数", 'perUnderBuild', lambda device, gpuName: {
'labelnames': ["建筑物下人群"],
'device': str(device),
'rainbows': COLOR,
'models': [
{
'weight': "../weights/trt/AIlib2/perUnderBuild/yolov5_%s_fp16.engine" % (gpuName), ###检测模型路径
'name': 'yolov5',
'model': yolov5Model,
'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.25, 'iou_thres': 0.45,
'segRegionCnt': 1, 'trtFlag_det': True, 'trtFlag_seg': False},
},
{
'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.50,
'gpu_id': 0,
'eos_coef': '0.5',
'set_cost_class': 1,
'set_cost_point': 0.05,
'backbone': 'vgg16_bn',
'expend': 10,
'psize': 5
},
}],
})
CITY_FIREAREA_MODEL = ("30", "307", "火焰面积模型", 'FireArea', lambda device, gpuName: {
'device': device,
'gpu_name': gpuName,
'labelnames': ["火焰"],
'seg_nclass': 2, # 分割模型类别数目默认2类
'segRegionCnt': 0,
'trtFlag_det': True,
'trtFlag_seg': False,
'Detweights': "../weights/trt/AIlib2/smogfire/yolov5_%s_fp16.engine" % gpuName, # 0fire 1smoke
'Samweights': "../weights/pth/AIlib2/firearea/sam_vit_b_01ec64.pth", #分割模型
'ksize':(7,7),
'sam_type':'vit_b',
'slopeIndex': [],
'segPar': None,
'postFile': {
"name": "post_process",
"conf_thres": 0.25,
"iou_thres": 0.45,
"classes": 5,
"rainbows": COLOR
},
'Segweights': None,
'fiterList':[1],
"score_byClass": {0: 0.1}
})
CITY_SECURITY_MODEL = ("30", "308", "安防模型", 'SECURITY', lambda device, gpuName: {
'labelnames': ["带安全帽","安全帽","攀爬","斗殴","未戴安全帽"],
'postProcess': {'function': security_post_process, 'pars': {'objs': [0,1],'iou':0.25,'unhelmet':4}},
'models':
[
{
'weight': "../weights/trt/AIlib2/security/yolov5_%s_fp16.engine" % (gpuName), ###检测模型路径
'name': 'yolov5',
'model': yolov5Model,
'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.25, 'iou_thres': 0.45,
'segRegionCnt': 1, 'trtFlag_det': True, 'trtFlag_seg': False},
}
],
'postFile': {
"rainbows": COLOR
},
'fiterList': [0,1],
"score_byClass": {"0": 0.50}
})
@staticmethod
def checkCode(code):
for model in ModelType:
if model.value[1] == code:
return True
return False
'''
参数1: 检测目标名称
参数2: 检测目标
参数3: 初始化百度检测客户端
'''
@unique
class BaiduModelTarget(Enum):
VEHICLE_DETECTION = (
"车辆检测", 0, lambda client0, client1, url, request_id: client0.vehicleDetectUrl(url, request_id))
HUMAN_DETECTION = (
"人体检测与属性识别", 1, lambda client0, client1, url, request_id: client1.bodyAttr(url, request_id))
PEOPLE_COUNTING = ("人流量统计", 2, lambda client0, client1, url, request_id: client1.bodyNum(url, request_id))
BAIDU_MODEL_TARGET_CONFIG = {
BaiduModelTarget.VEHICLE_DETECTION.value[1]: BaiduModelTarget.VEHICLE_DETECTION,
BaiduModelTarget.HUMAN_DETECTION.value[1]: BaiduModelTarget.HUMAN_DETECTION,
BaiduModelTarget.PEOPLE_COUNTING.value[1]: BaiduModelTarget.PEOPLE_COUNTING
}
EPIDEMIC_PREVENTION_CONFIG = {1: "行程码", 2: "健康码"}
# 模型分析方式
@unique
class ModelMethodTypeEnum(Enum):
# 方式一: 正常识别方式
NORMAL = 1
# 方式二: 追踪识别方式
TRACE = 2