75 lines
1.9 KiB
Python
75 lines
1.9 KiB
Python
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
|