gitadmin
/
tuoheng_algN
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

zsl #14

Merged
jiangchaoqing merged 7 commits from zsl into master 2025-08-25 11:20:50 +08:00
Conversation 0 Commits 7 Files Changed 11 +132 -81
5 changed files with 132 additions and 81 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit 2831d31f0b - Show all commits

2
concurrency/IntelligentRecognitionProcess.py
View File

@ -1051,6 +1051,8 @@ class PhotosIntelligentRecognitionProcess(Process):
ai_result_list = p_result[2]
for ai_result in ai_result_list:
box, score, cls = xywh2xyxy2(ai_result)
if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code):
box.append(ai_result[-1])
# 如果检测目标在识别任务中,继续处理
if cls in allowedList:
label_array = label_arraylist[cls]

21
concurrency/PushVideoStreamProcess.py
View File

@ -38,9 +38,8 @@ class PushStreamProcess(Process):
self._algStatus = False # 默认关闭
self._algSwitch = self._context['service']['algSwitch']
#0521:
default_enabled = str(self._msg.get("defaultEnabled", "True")).lower() == "true"
# 0521:
default_enabled = str(self._msg.get("defaultEnabled", "True")).lower() == "true"
if default_enabled:
print("执行默认程序（defaultEnabled=True）")
self._algSwitch = True
@ -168,6 +167,8 @@ class OnPushStreamProcess(PushStreamProcess):
else:
try: # 应对NaN情况
box, score, cls = xywh2xyxy2(qs)
if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code):
box.append(qs[-1])
except:
continue
if cls not in allowedList or score < frame_score:
@ -256,8 +257,10 @@ class OnPushStreamProcess(PushStreamProcess):
score = q[8]
rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"]
label_array, color = label_arrays[cls], rainbows[cls]
box = [(int(q[0]), int(q[1])), (int(q[2]), int(q[3])),
(int(q[4]), int(q[5])), (int(q[6]), int(q[7]))]
box = [(int(q[0]), int(q[1])), (int(q[2]), int(q[3])),
(int(q[4]), int(q[5])), (int(q[6]), int(q[7]))]
if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code):
box.append(qs[-1])
is_new = False
if q[11] == 1:
is_new = True
@ -412,6 +415,8 @@ class OffPushStreamProcess(PushStreamProcess):
else:
box, score, cls = xywh2xyxy2(qs)
if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code):
box.append(qs[-1])
if cls not in allowedList or score < frame_score:
continue
label_array, color = label_arrays[cls], rainbows[cls]
@ -494,8 +499,10 @@ class OffPushStreamProcess(PushStreamProcess):
score = q[8]
rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"]
label_array, color = label_arrays[cls], rainbows[cls]
box = [(int(q[0]), int(q[1])), (int(q[2]), int(q[3])),
(int(q[4]), int(q[5])), (int(q[6]), int(q[7]))]
box = [(int(q[0]), int(q[1])), (int(q[2]), int(q[3])),
(int(q[4]), int(q[5])), (int(q[6]), int(q[7]))]
if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code):
box.append(qs[-1])
is_new = False
if q[11] == 1:
is_new = True

98
enums/ModelTypeEnum.py
View File

@ -63,7 +63,7 @@ class ModelType(Enum):
"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
})
@ -99,10 +99,8 @@ class ModelType(Enum):
'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,'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 } },
'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False},
}
],
'postFile': {
@ -112,11 +110,10 @@ class ModelType(Enum):
"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] ],###控制哪些检测类别显示、输出
"score_byClass": {0: 0.25, 1: 0.3, 2: 0.3, 3: 0.3},
'fiterList': [5],
'segRegionCnt':2,###分割模型结果需要保留的等值线数目
"pixScale": 1.2,
})
@ -163,7 +160,7 @@ class ModelType(Enum):
"rainbows": COLOR
},
'score_byClass':{11:0.75,12:0.75},
'allowedList':[0,1,2,3,4,5,6,7,8,9,10,11,12,16,17,18,19,20,21,22],
'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
})
@ -232,7 +229,7 @@ class ModelType(Enum):
"classes": 5,
"rainbows": COLOR
},
'Segweights': None
'Segweights': None,
})
ANGLERSWIMMER_MODEL = ("9", "009", "钓鱼游泳模型", 'AnglerSwimmer', lambda device, gpuName: {
@ -378,9 +375,10 @@ 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,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 } }
'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,
@ -404,7 +402,7 @@ class ModelType(Enum):
"classes": 8,
"rainbows": COLOR
},
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6,7,8,9] ],###控制哪些检测类别显示、输出
"score_byClass":{0:0.8, 1:0.4, 2:0.5, 3:0.5},
'segRegionCnt':2,###分割模型结果需要保留的等值线数目
"pixScale": 1.2,
})
@ -569,10 +567,10 @@ class ModelType(Enum):
'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,'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} }
'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.1,'iou_thres':0.45,'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False}
},
{
'weight' : '../weights/pth/AIlib2/ocr2/crnn_ch.pth',
'weight' : '../weights/trt/AIlib2/ocr2/crnn_ch_%s_fp16_192X32.engine'%(gpuName),
'name':'ocr',
'model':ocrModel,
'par':{
@ -588,7 +586,6 @@ class ModelType(Enum):
},
}
],
'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6]],
'segPar': None,
'postFile': {
"name": "post_process",
@ -598,6 +595,8 @@ class ModelType(Enum):
"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: {
@ -643,7 +642,7 @@ class ModelType(Enum):
'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,'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 } },
'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False},
}
@ -656,7 +655,7 @@ class ModelType(Enum):
"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] ],###控制哪些检测类别显示、输出
"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,
@ -704,6 +703,7 @@ class ModelType(Enum):
"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
})
@ -717,7 +717,7 @@ class ModelType(Enum):
'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,'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 } },
'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'segRegionCnt':1, 'trtFlag_det':True,'trtFlag_seg':False},
}
@ -725,6 +725,7 @@ class ModelType(Enum):
'postFile': {
"rainbows": COLOR
},
"score_byClass": {0: 0.25, 1: 0.3, 2: 0.3, 3: 0.3}
})
@ -737,7 +738,7 @@ class ModelType(Enum):
'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,'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 } },
'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'segRegionCnt':1, 'trtFlag_det':True,'trtFlag_seg':False},
}
@ -745,6 +746,7 @@ class ModelType(Enum):
'postFile': {
"rainbows": COLOR
},
"score_byClass": {0: 0.25, 1: 0.3, 2: 0.3, 3: 0.3}
})
@ -757,7 +759,7 @@ class ModelType(Enum):
'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,'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 } },
'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'segRegionCnt':1, 'trtFlag_det':True,'trtFlag_seg':False },
}
@ -765,7 +767,6 @@ class ModelType(Enum):
'postFile': {
"rainbows": COLOR
},
})
TRAFFIC_SPILL_MODEL = ("50", "501", "高速公路抛洒物模型", 'highWaySpill', lambda device, gpuName: {
@ -807,7 +808,7 @@ class ModelType(Enum):
"classes": 2,
"rainbows": COLOR
},
'detModelpara': [{"id": str(x), "config": {"k1": "v1", "k2": "v2"}} for x in [0]],
'fiterList': [1],
###控制哪些检测类别显示、输出
'Detweights': "../weights/trt/AIlib2/highWaySpill/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/highWaySpill/stdc_360X640_%s_fp16.engine' % gpuName
@ -852,7 +853,7 @@ class ModelType(Enum):
"classes": 4,
"rainbows": COLOR
},
'detModelpara': [{"id": str(x), "config": {"k1": "v1", "k2": "v2"}} for x in [0]],
'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
@ -868,14 +869,15 @@ class ModelType(Enum):
'name': 'yolov5',
'model': yolov5Model,
'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.25, 'iou_thres': 0.45,
'allowedList': [0,1,2], 'segRegionCnt': 1, 'trtFlag_det': True,
'trtFlag_seg': False, "score_byClass": {"0": 0.25, "1": 0.3, "2": 0.3, "3": 0.3}},
'segRegionCnt': 1, 'trtFlag_det': True,
'trtFlag_seg': False},
}
],
'postFile': {
"rainbows": COLOR
},
'fiterList':[0]
})
@ -885,13 +887,13 @@ class ModelType(Enum):
'rainbows': COLOR,
'models': [
{
'trtFlag_det': False,
'weight': '../weights/pth/AIlib2/carplate/plate_yolov5s_v3.jit',
'name': 'yolov5',
'model': yolov5Model,
'par': {
'trtFlag_det': False,
'device': 'cuda:0',
'half': False,
'half': True,
'conf_thres': 0.4,
'iou_thres': 0.45,
'nc': 1,
@ -899,11 +901,11 @@ class ModelType(Enum):
},
},
{
'trtFlag_ocr': False,
'weight': '../weights/pth/AIlib2/ocr2/crnn_ch.pth',
'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,
@ -927,10 +929,8 @@ class ModelType(Enum):
'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}},
'segRegionCnt': 1, 'trtFlag_det': True,'trtFlag_seg': False},
}
],
'postFile': {
"rainbows": COLOR
@ -948,8 +948,7 @@ class ModelType(Enum):
'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}},
'segRegionCnt': 1, 'trtFlag_det': True, 'trtFlag_seg': False},
}
],
@ -996,8 +995,7 @@ class ModelType(Enum):
'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}},
'segRegionCnt': 1, 'trtFlag_det': True, 'trtFlag_seg': False},
}
],
@ -1017,8 +1015,7 @@ class ModelType(Enum):
'name': 'yolov5',
'model': yolov5Model,
'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.25, 'iou_thres': 0.45,
'allowedList': [0,1,2], 'segRegionCnt': 1, 'trtFlag_det': True,
'trtFlag_seg': False, "score_byClass": {"0": 0.25, "1": 0.3, "2": 0.3, "3": 0.3}},
'segRegionCnt': 1, 'trtFlag_det': True, 'trtFlag_seg': False},
},
{
'trtFlag_det': False,
@ -1043,6 +1040,33 @@ class ModelType(Enum):
}],
})
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, # 0：fire 1：smoke
'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}
})
@staticmethod
def checkCode(code):
for model in ModelType:

68
util/ModelUtils.py
View File

@ -27,6 +27,7 @@ import torch
import tensorrt as trt
from utilsK.jkmUtils import pre_process, post_process, get_return_data
from DMPR import DMPRModel
from segment_anything import SamPredictor, sam_model_registry
FONT_PATH = "../AIlib2/conf/platech.ttf"
@ -35,8 +36,8 @@ class OneModel:
__slots__ = "model_conf"
def __init__(self, device, allowedList=None, requestId=None, modeType=None, gpu_name=None, base_dir=None, env=None):
s = time.time()
try:
start = time.time()
logger.info("########################加载{}########################, requestId:{}", modeType.value[2],
requestId)
par = modeType.value[4](str(device), gpu_name)
@ -69,11 +70,11 @@ class OneModel:
'ovlap_thres_crossCategory': postFile.get("ovlap_thres_crossCategory"),
'iou_thres': postFile["iou_thres"],
# 对高速模型进行过滤
'allowedList': par['allowedList'] if modeType.value[0] == '3' else [],
'segRegionCnt': par['segRegionCnt'],
'trtFlag_det': par['trtFlag_det'],
'trtFlag_seg': par['trtFlag_seg'],
'score_byClass':par['score_byClass'] if 'score_byClass' in par.keys() else None
'score_byClass':par['score_byClass'] if 'score_byClass' in par.keys() else None,
'fiterList': par['fiterList'] if 'fiterList' in par.keys() else []
}
model_param = {
"model": model,
@ -88,7 +89,7 @@ class OneModel:
logger.error("模型加载异常：{}, requestId:{}", format_exc(), requestId)
raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
logger.info("模型初始化时间：{}, requestId:{}", time.time() - s, requestId)
logger.info("模型初始化时间：{}, requestId:{}", time.time() - start, requestId)
# 纯分类模型
class cityManagementModel:
__slots__ = "model_conf"
@ -106,6 +107,8 @@ class cityManagementModel:
model_param = {
"modelList": modelList,
"postProcess": postProcess,
"score_byClass":par['score_byClass'] if 'score_byClass' in par.keys() else None,
"fiterList":par['fiterList'] if 'fiterList' in par.keys() else [],
}
self.model_conf = (modeType, model_param, allowedList, names, rainbows)
except Exception:
@ -114,15 +117,14 @@ class cityManagementModel:
ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
def detSeg_demo2(args):
model_conf, frame, request_id = args
modelList, postProcess = model_conf[1]['modelList'], model_conf[1]['postProcess']
modelList, postProcess,score_byClass,fiterList = (
model_conf[1]['modelList'], model_conf[1]['postProcess'],model_conf[1]['score_byClass'], model_conf[1]['fiterList'])
try:
result = [[ None, None, AI_process_N([frame], modelList, postProcess)[0] ] ] # 为了让返回值适配统一的接口而写的shi
result = [[ None, None, AI_process_N([frame], modelList, postProcess,score_byClass,fiterList)[0] ] ] # 为了让返回值适配统一的接口而写的shi
return result
except ServiceException as s:
raise s
except Exception:
# self.num += 1
# cv2.imwrite('/home/th/tuo_heng/dev/img%s.jpg' % str(self.num), frame)
logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), request_id)
raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0],
ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1])
@ -130,11 +132,6 @@ def detSeg_demo2(args):
def model_process(args):
model_conf, frame, request_id = args
model_param, names, rainbows = model_conf[1], model_conf[3], model_conf[4]
# modeType, model_param, allowedList, names, rainbows = model_conf
# segmodel, names, label_arraylist, rainbows, objectPar, font, segPar, mode, postPar, requestId = args
# model_param['digitFont'] = digitFont
# model_param['label_arraylist'] = label_arraylist
# model_param['font_config'] = font_config
try:
return AI_process([frame], model_param['model'], model_param['segmodel'], names, model_param['label_arraylist'],
rainbows, objectPar=model_param['objectPar'], font=model_param['digitFont'],
@ -167,7 +164,13 @@ class TwoModel:
Detweights = par['Detweights']
with open(Detweights, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.ERROR)) as runtime:
model = runtime.deserialize_cuda_engine(f.read())
segmodel = None
if modeType == ModelType.CITY_FIREAREA_MODEL:
sam = sam_model_registry[par['sam_type']](checkpoint=par['Samweights'])
sam.to(device=device)
segmodel = SamPredictor(sam)
else:
segmodel = None
postFile = par['postFile']
conf_thres = postFile["conf_thres"]
iou_thres = postFile["iou_thres"]
@ -181,7 +184,10 @@ class TwoModel:
"conf_thres": conf_thres,
"iou_thres": iou_thres,
"trtFlag_det": par['trtFlag_det'],
"otc": otc
"otc": otc,
"ksize":par['ksize'] if 'ksize' in par.keys() else None,
"score_byClass": par['score_byClass'] if 'score_byClass' in par.keys() else None,
"fiterList": par['fiterList'] if 'fiterList' in par.keys() else []
}
self.model_conf = (modeType, model_param, allowedList, names, rainbows)
except Exception:
@ -189,16 +195,15 @@ class TwoModel:
raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
logger.info("模型初始化时间：{}, requestId:{}", time.time() - s, requestId)
def forest_process(args):
model_conf, frame, request_id = args
model_param, names, rainbows = model_conf[1], model_conf[3], model_conf[4]
try:
return AI_process_forest([frame], model_param['model'], model_param['segmodel'], names,
model_param['label_arraylist'], rainbows, model_param['half'], model_param['device'],
model_param['conf_thres'], model_param['iou_thres'], [], font=model_param['digitFont'],
trtFlag_det=model_param['trtFlag_det'], SecNms=model_param['otc'])
model_param['conf_thres'], model_param['iou_thres'],font=model_param['digitFont'],
trtFlag_det=model_param['trtFlag_det'], SecNms=model_param['otc'],ksize = model_param['ksize'],
score_byClass=model_param['score_byClass'],fiterList=model_param['fiterList'])
except ServiceException as s:
raise s
except Exception:
@ -207,7 +212,6 @@ def forest_process(args):
logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), request_id)
raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0],
ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1])
class MultiModel:
__slots__ = "model_conf"
@ -226,6 +230,8 @@ class MultiModel:
model_param = {
"modelList": modelList,
"postProcess": postProcess,
"score_byClass": par['score_byClass'] if 'score_byClass' in par.keys() else None,
"fiterList": par['fiterList'] if 'fiterList' in par.keys() else []
}
self.model_conf = (modeType, model_param, allowedList, names, rainbows)
except Exception:
@ -233,13 +239,13 @@ class MultiModel:
raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
logger.info("模型初始化时间：{}, requestId:{}", time.time() - s, requestId)
def channel2_process(args):
model_conf, frame, request_id = args
modelList, postProcess = model_conf[1]['modelList'], model_conf[1]['postProcess']
modelList, postProcess,score_byClass,fiterList = (
model_conf[1]['modelList'], model_conf[1]['postProcess'],model_conf[1]['score_byClass'], model_conf[1]['fiterList'])
try:
start = time.time()
result = [[None, None, AI_process_C([frame], modelList, postProcess)[0]]] # 为了让返回值适配统一的接口而写的shi
result = [[None, None, AI_process_C([frame], modelList, postProcess,score_byClass,fiterList)[0]]] # 为了让返回值适配统一的接口而写的shi
# print("AI_process_C use time = {}".format(time.time()-start))
return result
except ServiceException as s:
@ -248,7 +254,6 @@ def channel2_process(args):
logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), request_id)
raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0],
ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1])
def get_label_arraylist(*args):
width, height, names, rainbows = args
# line = int(round(0.002 * (height + width) / 2) + 1)
@ -269,8 +274,6 @@ def get_label_arraylist(*args):
'label_location': 'leftTop'}
label_arraylist = get_label_arrays(names, rainbows, fontSize=text_height, fontPath=FONT_PATH)
return digitFont, label_arraylist, (line, text_width, text_height, fontScale, tf)
# 船只模型
class ShipModel:
__slots__ = "model_conf"
@ -296,8 +299,6 @@ class ShipModel:
raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
logger.info("模型初始化时间：{}, requestId:{}", time.time() - s, requestId)
def obb_process(args):
model_conf, frame, request_id = args
model_param = model_conf[1]
@ -312,7 +313,6 @@ def obb_process(args):
logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), request_id)
raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0],
ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1])
# 车牌分割模型、健康码、行程码分割模型
class IMModel:
__slots__ = "model_conf"
@ -336,7 +336,7 @@ class IMModel:
new_device = torch.device(par['device'])
model = torch.jit.load(par[img_type]['weights'])
logger.info("########################加载 jit 模型成功 成功 ########################, requestId:{}",
logger.info("########################加载 jit 模型成功 成功 ########################, requestId:{}",
requestId)
self.model_conf = (modeType, allowedList, new_device, model, par, img_type)
except Exception:
@ -398,7 +398,6 @@ class CARPLATEModel:
raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
class DENSECROWDCOUNTModel:
__slots__ = "model_conf"
@ -766,4 +765,11 @@ MODEL_CONFIG = {
None,
lambda x: cc_process(x)
),
# 加载火焰面积模型
ModelType.CITY_FIREAREA_MODEL.value[1]: (
lambda x, y, r, t, z, h: TwoModel(x, y, r, ModelType.CITY_FIREAREA_MODEL, t, z, h),
ModelType.CITY_FIREAREA_MODEL,
lambda x, y, z: one_label(x, y, z),
lambda x: forest_process(x)
),
}

24
util/PlotsUtils.py
View File

@ -77,10 +77,19 @@ def xy2xyxy(box):
def draw_painting_joint(box, img, label_array, score=0.5, color=None, config=None, isNew=False):
# 识别问题描述图片的高、宽
lh, lw = label_array.shape[0:2]
# 图片的长度和宽度
imh, imw = img.shape[0:2]
box = xywh2xyxy(box)
points = np.array([])
if isinstance(box[-1], np.ndarray):
font = ImageFont.truetype(FONT_PATH, 10, encoding='utf-8')
points = box[-1]
arrea = cv2.contourArea(points)
label = '火焰面积:%s'%arrea
label_array = get_label_array(color, label, font, 10)
else:
label = ' %.2f' % score
lh, lw = label_array.shape[0:2]
box = xywh2xyxy(box[:4])
# 框框左上的位置
x0, y1 = box[0][0], box[0][1]
# if score_location == 'leftTop':
@ -121,13 +130,16 @@ def draw_painting_joint(box, img, label_array, score=0.5, color=None, config=Non
6. shift：顶点坐标中小数的位数
'''
tl = config[0]
box1 = np.asarray(box, np.int32)
cv2.polylines(img, [box1], True, color, tl)
img[y0:y1, x0:x1, :] = label_array
if points.size != 0:
cv2.drawContours(img, points, -1, color, 2)
return img, box
else:
box1 = np.asarray(box, np.int32)
cv2.polylines(img, [box1], True, color, tl)
pts_cls = [(x0, y0), (x1, y1)]
# 把英文字符score画到类别旁边
# tl = max(int(round(imw / 1920 * 3)), 1) or round(0.002 * (imh + imw) / 2) + 1
label = ' %.2f' % score
# tf = max(tl, 1)
# fontScale = float(format(imw / 1920 * 1.1, '.2f')) or tl * 0.33
# fontScale = tl * 0.33
@ -230,7 +242,6 @@ def draw_name_ocr(box, img, color, line_thickness=2, outfontsize=40):
# (color=None, label=None, font=None, fontSize=40, unify=False)
label_zh = get_label_array(color, box[0], font, outfontsize)
return plot_one_box_auto(box[1], img, color, line_thickness, label_zh)
def filterBox(det0, det1, pix_dis):
# det0为 (m1, 11) 矩阵
# det1为 (m2, 12) 矩阵
@ -276,6 +287,7 @@ def plot_one_box_auto(box, img, color=None, line_thickness=2, label_array=None):
# print("省略 :%s, lh:%s, lw:%s"%('+++' * 10, lh, lw))
# 图片的长度和宽度
imh, imw = img.shape[0:2]
points = None
box = xy2xyxy(box)
# 框框左上的位置
x0, y1 = box[0][0], box[0][1]