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
from utilsK.drownUtils import mixDrowing_water_postprocess
from utilsK.noParkingUtils import mixNoParking_road_postprocess
from utilsK.illParkingUtils import illParking_postprocess
from stdc import stdcModel
from yolov5 import yolov5Model
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
        },
        'Detweights': "../AIlib2/weights/river/yolov5_%s_fp16.engine" % gpuName,
        'Segweights': '../AIlib2/weights/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': "../AIlib2/weights/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':"../AIlib2/weights/forest2/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.3,"2":0.3,"3":0.3  } },       
                    }
                    
                        
                ],
                
        'postFile': {
            "name": "post_process",
            "conf_thres": 0.25,
            "iou_thres": 0.45,
            "classes": 5,
            "rainbows": COLOR
        },
        '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,

 
    })
     


    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),
                    #'modelSize': (1920,1080),
                    'RoadArea': 16000,
                    'roadVehicleAngle': 15,
                    'speedRoadVehicleAngleMax': 75,
                    'roundness': 1.0,
                    'cls': 9,
                    '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
        },
        'Detweights': "../AIlib2/weights/highWay2/yolov5_%s_fp16.engine" % gpuName,
        'Segweights': '../AIlib2/weights/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': "../AIlib2/weights/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': "../AIlib2/weights/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': "../AIlib2/weights/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': "../AIlib2/weights/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': "../AIlib2/weights/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': '../AIlib2/weights/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': "../AIlib2/weights/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
                }
            }
        },
        'postFile': {
            "name": "post_process",
            "conf_thres": 0.25,
            "iou_thres": 0.3,
            "ovlap_thres_crossCategory": 0.65,
            "classes": 5,
            "rainbows": COLOR
        },
        # "../AIlib2/weights/conf/%s/yolov5.pt" % modeType.value[3]
        'Detweights': "../AIlib2/weights/river2/yolov5_%s_fp16.engine" % gpuName,
        # '../AIlib2/weights/conf/%s/stdc_360X640.pth' % modeType.value[3]
        'Segweights': '../AIlib2/weights/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':3,'scaleRatio':0.5,'border':80}
                    },
        'models':[
            {
                #'weight':'../AIlib2/weights/conf/cityMangement3/yolov5.pt',
                'weight':'../AIlib2/weights/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.8,"1":0.4,"2":0.5,"3":0.5  } }
            },
            {
                'weight':'../AIlib2/weights/conf/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':'../AIlib2/weights/conf/cityMangement3/stdc_360X640.pth',
                '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},###分割模型预处理参数
                'model':stdcModel,
                'name':'stdc'           
            } 
        ],
        'postFile': {
            "name": "post_process",
            "conf_thres": 0.25,
            "iou_thres": 0.45,
            "classes": 5,
            "rainbows": COLOR
        },
        'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9]  ],###控制哪些检测类别显示、输出
        '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
        },
        # "../AIlib2/weights/conf/%s/yolov5.pt" % modeType.value[3]
        'Detweights': "../AIlib2/weights/drowning/yolov5_%s_fp16.engine" % gpuName,
        # '../AIlib2/weights/conf/%s/stdc_360X640.pth' % modeType.value[3]
        'Segweights': '../AIlib2/weights/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': "../AIlib2/weights/noParking/yolov5_%s_fp16.engine" % gpuName,
            'Segweights': '../AIlib2/weights/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': "../AIlib2/weights/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': "../AIlib2/weights/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': "../AIlib2/weights/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': ["国旗", "浮标", "船名", "船只","未挂国旗船只"],
        'segRegionCnt': 0,
        'postProcess':{'function':channel2_post_process,'name':'channel2','pars':{
            'objs':[2],
            'wRation':1/6.0,
            'hRation':1/6.0,
            'smallId':0,
            'bigId':3,
            'newId':4,
            'recScale':1.2}},       
        'models':[
        {
            #'weight':'../AIlib2/weights/conf/channel2/yolov5.pt',
            'weight':'../AIlib2/weights/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' : '../AIlib2/weights/ocr2/crnn_ch_4090_fp16_192X32.engine',    
            'weight' : '../AIlib2/weights/conf/ocr2/crnn_ch.pth',
            'name':'ocr',
            'model':ocrModel,
            'par':{
                'char_file':'../AIlib2/weights/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,
                },
        }  
        ],
        'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3]],
        'segPar': None,
        'postFile': {
            "name": "post_process",
            "conf_thres": 0.25,
            "iou_thres": 0.45,
            "classes": 5,
            "rainbows": COLOR
        },
        'Segweights': None,
    })
    
    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
        },
        # "../AIlib2/weights/conf/%s/yolov5.pt" % modeType.value[3]
        'Detweights': "../AIlib2/weights/riverT/yolov5_%s_fp16.engine" % gpuName,
        # '../AIlib2/weights/conf/%s/stdc_360X640.pth' % modeType.value[3]
        'Segweights': '../AIlib2/weights/riverT/stdc_360X640_%s_fp16.engine' % gpuName
    })



    FORESTCROWD_FARM_MODEL = ("2", "026", "森林人群模型", 'forestCrowd', lambda device, gpuName: {
        'labelnames': ["林斑", "病死树", "行人", "火焰", "烟雾","人群"],
        'postProcess':{'function':gather_post_process,'pars':{'pedestrianId':2,'crowdThreshold':4,'gatherId':5,'distancePersonScale':2.0}},
        'models':
                [
                    {
                        'weight':"../AIlib2/weights/forestCrowd/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.25,"2":0.6,"3":0.6,'4':0.6 ,'5':0.6   } },
                    }


                ],

        'postFile': {
            "name": "post_process",
            "conf_thres": 0.25,
            "iou_thres": 0.45,
            "classes": 5,
            "rainbows": COLOR
        },
        '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,


    })
    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),
                    #'modelSize': (1920,1080),
                    'RoadArea': 16000,
                    'roadVehicleAngle': 15,
                    'speedRoadVehicleAngleMax': 75,
                    'roundness': 1.0,
                    'cls': 9,
                    '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
        },
        'Detweights': "../AIlib2/weights/highWay2/yolov5_%s_fp16.engine" % gpuName,
        'Segweights': '../AIlib2/weights/highWay2/stdc_360X640_%s_fp16.engine' % gpuName
    })



    @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
