import sys
from enum import Enum, unique

from common.Constant import COLOR

sys.path.extend(['..', '../AIlib2'])
from utilsK.queRiver import riverDetSegMixProcess

'''
参数说明
1. 编号  
2. 模型编号
3. 模型名称
4. 选用的模型名称
5. 模型配置
6. 模型引用配置[Detweights文件, Segweights文件, 引用计数]
'''


@unique
class ModelType(Enum):


    RIVER2_MODEL = ("15", "015", "河道检测模型", 'river2', lambda device, gpuName: {
        # 实际具体服务器显卡数
        'device': '0,1',
        'labelnames': ["漂浮物", "岸坡垃圾", "排口", "违建", "菜地", "水生植物", "河湖人员", "钓鱼人员", "船只",
                       "蓝藻"],
        'trtFlag_seg': False,
        'trtFlag_det': False,
        '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/pth/AIlib2/river2/yolov5.pt",
        'Segweights': '../weights/pth/AIlib2/river2/stdc_360X640.pth'
    })

    @staticmethod
    def checkCode(code):
        for model in ModelType:
            if model.value[1] == code:
                return True
        return False


# 模型分析方式
@unique
class ModelMethodTypeEnum(Enum):
    # 方式一: 正常识别方式
    NORMAL = 1

    # 方式二: 追踪识别方式
    TRACE = 2