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

Compare commits

base: gitadmin:dsj
Branches Tags
gitadmin:master
gitadmin:zsl
gitadmin:dsj
gitadmin:bazong
gitadmin:jcq
gitadmin:wj
..
compare: gitadmin:master
Branches Tags
gitadmin:master
gitadmin:zsl
gitadmin:dsj
gitadmin:bazong
gitadmin:jcq
gitadmin:wj

17 Commits
dsj ... master

Author SHA1 Message Date
jiangchaoqing 79a859e886 Merge pull request 'zsl' (#14) from zsl into master 
Reviewed-on: #14
 2025-08-25 11:20:27 +08:00
jiangchaoqing b8fc17c39f revert e5e778cf3c 
revert update
 2025-08-25 11:15:46 +08:00
th bea36e2601 1）自研车牌模型支持engine 2025-08-23 17:08:47 +08:00
th e9098d26d2 1）新增区域入侵算法:mqtt服务不开启 2）城管dmpr支撑engine 2025-08-23 09:14:44 +08:00
th 0f3a0a85b2 1）新增区域入侵算法 2）火焰面积bug修复 2025-08-22 11:14:08 +08:00
th 6e89b6587b 1)火焰面积模型bug修复 2）新增安防模型 2025-08-16 11:15:41 +08:00
th e82f07aa1a riverDetSegMixProcess:支持指定河边（河道膨胀指定scale范围类）过滤 2025-08-11 14:32:23 +08:00
jiangchaoqing e5e778cf3c update 2025-08-11 10:27:12 +08:00
th 2831d31f0b 1)新增M029 火焰面积 2）算法支持按类过滤 3）算法支持按置信度过滤 4）其他优化 2025-08-09 17:44:33 +08:00
th 2d626929f7 model_process:支持按类置信度筛选类别 2025-07-31 19:33:00 +08:00
jiangchaoqing 2c4b880d93 Merge pull request 'zsl' (#8) from zsl into master 
Reviewed-on: #8
 2025-07-26 14:23:29 +08:00
th b46182d68e 优化人群计数人数显示及边界设置 2025-07-26 10:05:31 +08:00
th 585a7a05f7 1)新增M027:建筑物下行人检测及计数 2)密集人群计数模型及自研车牌模型优化 3)分类模型支持pt模型加载（巴中水利） 2025-07-25 18:52:10 +08:00
th b75a74d52c 1)新增M027:建筑物下行人检测及计数 2)密集人群计数模型及自研车牌模型优化 3)分类模型支持pt模型加载（巴中水利） 2025-07-25 18:47:32 +08:00
th 98480b45d6 1)新增M027:建筑物下行人检测及计数 2)密集人群计数模型及自研车牌模型优化 3)分类模型支持pt模型加载（巴中水利） 2025-07-25 18:25:50 +08:00
th d369031085 1)新增M027:建筑物下行人检测及计数 2)密集人群计数模型及自研车牌模型优化 3)分类模型支持pt模型加载（巴中水利） 2025-07-25 18:20:18 +08:00
th 12a4b296e1 车牌及健康码权重文件路径优化 2025-06-26 13:28:04 +08:00
14 changed files with 2517 additions and 666 deletions
Show Stats Expand all files Collapse all files

10
README.md
View File

@ -1,2 +1,8 @@
1.dsj 1.2025.01.21把之前的tuoheng alg仓库代码重新开个仓库 (1)在config/service/dsp_test_service.yml里面添加参数，控制存储用的oss还是minio storage_source: 1 2.2025.02.06 (1)修改代码，把mqtt读取加入到系统中。config/service/dsp_test_service.yml，中添加mqtt_flag,决定是否启用。 (2)修改了minio情况下的，文件名命名方式。 3.2025.02.12 (1)增加了对alg算法开发的代码。可以通过配置文件config/service/dsp_test_service.yml中algSwitch: true，决定是否启用。
2.0811
4、2025.07.10 周树亮 - 增加人群计数，自研车牌模型，裸土覆盖3个场景
5、江朝庆 -- 0715
1）代码整理，删除冗余代码。
2）增加requirements.txt,方便部署
3) logs

4
common/Constant.py
View File

@ -16,7 +16,6 @@ success_progess = "1.0000"
width = 1400 width = 1400
COLOR = ( COLOR = (
[0, 0, 255],
[255, 0, 0], [255, 0, 0],
[211, 0, 148], [211, 0, 148],
[0, 127, 0], [0, 127, 0],
@ -35,7 +34,8 @@ COLOR = (
[8, 101, 139], [8, 101, 139],
[171, 130, 255], [171, 130, 255],
[139, 112, 74], [139, 112, 74],
[205, 205, 180]) [205, 205, 180],
[0, 0, 255],)
ONLINE = "online" ONLINE = "online"
OFFLINE = "offline" OFFLINE = "offline"

366
concurrency/FileUploadThread.py
View File

@ -3,52 +3,38 @@ from concurrent.futures import ThreadPoolExecutor
from threading import Thread from threading import Thread
from time import sleep, time from time import sleep, time
from traceback import format_exc from traceback import format_exc
import numpy as np
from loguru import logger from loguru import logger
import cv2 import cv2
import requests
import json
import base64
from entity.FeedBack import message_feedback from entity.FeedBack import message_feedback
from enums.ExceptionEnum import ExceptionType from enums.ExceptionEnum import ExceptionType
from enums.ModelTypeEnum import ModelType
from exception.CustomerException import ServiceException from exception.CustomerException import ServiceException
from util.AliyunSdk import AliyunOssSdk from util.AliyunSdk import AliyunOssSdk
from util.MinioSdk import MinioSdk from util.MinioSdk import MinioSdk
from util import TimeUtils from util import TimeUtils
from enums.AnalysisStatusEnum import AnalysisStatus from enums.AnalysisStatusEnum import AnalysisStatus
from util.PlotsUtils import draw_painting_joint from util.PlotsUtils import draw_painting_joint, draw_name_ocr, draw_name_crowd,draw_transparent_red_polygon
from util.QueUtil import put_queue, get_no_block_queue, clear_queue from util.QueUtil import put_queue, get_no_block_queue, clear_queue
import io import io
from util.LocationUtils import locate_byMqtt from util.LocationUtils import locate_byMqtt
class FileUpload(Thread): class FileUpload(Thread):
__slots__ = ('_fb_queue', '_context', '_image_queue', '_analyse_type', '_msg', '_mqtt_list') __slots__ = ('_fb_queue', '_context', '_image_queue', '_analyse_type', '_msg')
def __init__(self, *args): def __init__(self, *args):
super().__init__() super().__init__()
self._fb_queue, self._context, self._msg, self._image_queue, self._analyse_type,self._mqtt_list = args self._fb_queue, self._context, self._msg, self._image_queue, self._analyse_type = args
self._storage_source = self._context['service']['storage_source'] self._storage_source = self._context['service']['storage_source']
self._algStatus = False # 默认关闭 self._algStatus = False # 默认关闭
# self._algStatus = True # 默认关闭
self._algSwitch = self._context['service']['algSwitch']
# self._algStatus = True # 默认关闭
self.server_ip = self._context['llm']['server_ip'] self._algSwitch = self._context['service']['algSwitch']
self.server_port = self._context['llm']['server_port']
self.is_use_llm = self._context['llm']['is_use_llm'] # 0521:
self.cnt = self._context['llm']['cnt'] default_enabled = str(self._msg.get("defaultEnabled", "True")).lower() == "true"
#0521:
default_enabled = str(self._msg.get("defaultEnabled", "True")).lower() == "true"
if default_enabled: if default_enabled:
print("执行默认程序（defaultEnabled=True）") print("执行默认程序（defaultEnabled=True）")
self._algSwitch = True self._algSwitch = True
@ -58,73 +44,14 @@ class FileUpload(Thread):
# 这里放非默认逻辑的代码 # 这里放非默认逻辑的代码
self._algSwitch = False self._algSwitch = False
print("---line46 :FileUploadThread.py---", self._algSwitch)
# 如果任务是在线、离线处理，则用此类 # 如果任务是在线、离线处理，则用此类
class ImageFileUpload(FileUpload): class ImageFileUpload(FileUpload):
__slots__ = () __slots__ = ()
def extract_answers(self, text): # 提取大模型中的分点文本内容 # @staticmethod
lines = [line.strip() for line in text.split('.') if line.strip()]
answers = []
for i, line in enumerate(lines):
if i > 0:
if i < len(lines) - 1:
answers.append(line[:-1])
else:
answers.append(line)
return answers
def get_llm_res(self, image, prompt,SERVER_IP,API_PORT):
# SERVER_IP = "192.168.10.11"
# API_URL = f"http://{SERVER_IP}:8000/generate"
SERVER_IP = SERVER_IP
API_URL = f"http://{SERVER_IP}:{API_PORT}/generate"
# _, buffer = cv2.imencode('.jpg', image)
# is_success, buffer = cv2.imencode('.png', image)
is_success, buffer = cv2.imencode('.webp', image, [cv2.IMWRITE_WEBP_QUALITY, 100])
image_base64 = base64.b64encode(buffer).decode('utf-8')
conversation = {
"prompt": prompt,
"image_base64": image_base64,
}
try:
logger.info(f"正在向 {API_URL} 发送请求...")
response = requests.post(API_URL, json=conversation, timeout=60) # 设置超时
response.raise_for_status() # 如果状态码不是2xx，会抛出异常
resp = response.json()
result = resp.get("response")
logger.info("\n--- 请求成功 ---")
logger.info(f"原始Prompt: {prompt}")
logger.info(f"模型生成结果: {result}")
except requests.exceptions.RequestException as e:
logger.info(f"\n--- 请求失败 ---")
logger.info(f"发生错误: {e}")
except json.JSONDecodeError:
logger.info("\n--- 解析失败 ---")
logger.info(f"无法解析服务器返回的响应: {response.text}")
new_result = self.extract_answers(result)
print("###line116",new_result)
new_result = new_result[1] # 获取第二点内容
# flag = "不" in result or ("没" in result and "没错" not in result) or "否" in result or "未" in result or "无" in result
flag = "不存在" in new_result or ("没" in new_result and "没错" not in new_result) or (
"否" in new_result and "是否" not in new_result) or "未" in new_result or "无法" in new_result
return not flag, new_result
#@staticmethod
def handle_image(self, frame_msg, frame_step): def handle_image(self, frame_msg, frame_step):
# (high_score_image["code"], all_frames, draw_config["font_config"]) # (high_score_image["code"], all_frames, draw_config["font_config"])
# high_score_image["code"][code][cls] = (frame, frame_index_list[i], cls_list) # high_score_image["code"][code][cls] = (frame, frame_index_list[i], cls_list)
@ -138,25 +65,40 @@ class ImageFileUpload(FileUpload):
模型编号：modeCode 模型编号：modeCode
检测目标：detectTargetCode 检测目标：detectTargetCode
''' '''
print('*'*100,' mqtt_list:',len(self._mqtt_list)) aFrame = frame.copy()
igH, igW = aFrame.shape[0:2]
model_info = [] model_info = []
mqttPares= det_xywh['mqttPares']
border = None
gps = [None, None]
camParas = None
if mqttPares is not None:
if mqttPares[0] == 1:
border = mqttPares[1]
elif mqttPares[0] == 0:
camParas = mqttPares[1]
if border is not None:
aFrame = draw_transparent_red_polygon(aFrame, np.array(border, np.int32), alpha=0.25)
det_xywh.pop('mqttPares')
# 更加模型编码解析数据 # 更加模型编码解析数据
for code, det_list in det_xywh.items(): for code, det_list in det_xywh.items():
if len(det_list) > 0: if len(det_list) > 0:
for cls, target_list in det_list.items(): for cls, target_list in det_list.items():
if len(target_list) > 0: if len(target_list) > 0:
aFrame = frame.copy()
for target in target_list: for target in target_list:
draw_painting_joint(target[1], aFrame, target[3], target[2], target[4], font_config, target[5]) if camParas is not None:
igH,igW = aFrame.shape[0:2] gps = locate_byMqtt(target[1], igW, igH, camParas, outFormat='wgs84')
if len(self._mqtt_list)>=1: # 自研车牌模型判断
#camParas = self._mqtt_list[0]['data'] if ModelType.CITY_CARPLATE_MODEL.value[1] == str(code):
camParas = self._mqtt_list[0] draw_name_ocr(target[1], aFrame, target[4])
gps = locate_byMqtt(target[1],igW,igH,camParas,outFormat='wgs84') elif ModelType.CITY_DENSECROWDCOUNT_MODEL.value[1] == str(code) or\
else: ModelType.CITY_UNDERBUILDCOUNT_MODEL.value[1] == str(code):
gps=[None,None] draw_name_crowd(target[1], aFrame, target[4])
model_info.append({"modelCode": str(code), "detectTargetCode": str(cls), "aFrame": aFrame,'gps':gps}) else:
draw_painting_joint(target[1], aFrame, target[3], target[2], target[4], font_config,
target[5],border)
model_info.append(
{"modelCode": str(code), "detectTargetCode": str(cls), "aFrame": aFrame, 'gps': gps})
if len(model_info) > 0: if len(model_info) > 0:
image_result = { image_result = {
"or_frame": frame, "or_frame": frame,
@ -175,13 +117,15 @@ class ImageFileUpload(FileUpload):
image_queue, fb_queue, analyse_type = self._image_queue, self._fb_queue, self._analyse_type image_queue, fb_queue, analyse_type = self._image_queue, self._fb_queue, self._analyse_type
service_timeout = int(service["timeout"]) service_timeout = int(service["timeout"])
frame_step = int(service["filter"]["frame_step"]) + 120 frame_step = int(service["filter"]["frame_step"]) + 120
if msg['taskType']==0: self._algStatus = False if msg['taskType'] == 0:
else: self._algStatus = True self._algStatus = False
else:
self._algStatus = True
try: try:
with ThreadPoolExecutor(max_workers=2) as t: with ThreadPoolExecutor(max_workers=2) as t:
# 初始化oss客户端 # 初始化oss客户端
if self._storage_source==1: if self._storage_source == 1:
minioSdk = MinioSdk(base_dir, env, request_id ) minioSdk = MinioSdk(base_dir, env, request_id)
else: else:
aliyunOssSdk = AliyunOssSdk(base_dir, env, request_id) aliyunOssSdk = AliyunOssSdk(base_dir, env, request_id)
start_time = time() start_time = time()
@ -195,15 +139,16 @@ class ImageFileUpload(FileUpload):
# 获取队列中的消息 # 获取队列中的消息
image_msg = get_no_block_queue(image_queue) image_msg = get_no_block_queue(image_queue)
if image_msg is not None: if image_msg is not None:
if image_msg[0] == 2: if image_msg[0] == 2:
logger.info("图片上传线程收到命令:{}, requestId: {}",image_msg[1] ,request_id) logger.info("图片上传线程收到命令:{}, requestId: {}", image_msg[1], request_id)
if 'stop' == image_msg[1]: if 'stop' == image_msg[1]:
logger.info("开始停止图片上传线程, requestId:{}", request_id) logger.info("开始停止图片上传线程, requestId:{}", request_id)
break break
if 'algStart' == image_msg[1]: self._algStatus = True; logger.info("图片上传线程,执行算法开启命令, requestId:{}", request_id) if 'algStart' == image_msg[1]: self._algStatus = True; logger.info(
if 'algStop' == image_msg[1]: self._algStatus = False; logger.info("图片上传线程,执行算法关闭命令, requestId:{}", request_id) "图片上传线程,执行算法开启命令, requestId:{}", request_id)
if 'algStop' == image_msg[1]: self._algStatus = False; logger.info(
"图片上传线程,执行算法关闭命令, requestId:{}", request_id)
if image_msg[0] == 1: if image_msg[0] == 1:
image_result = self.handle_image(image_msg[1], frame_step) image_result = self.handle_image(image_msg[1], frame_step)
if image_result is not None: if image_result is not None:
@ -213,117 +158,54 @@ class ImageFileUpload(FileUpload):
image_result["last_frame"], image_result["last_frame"],
analyse_type, analyse_type,
"OR", "0", "0", request_id) "OR", "0", "0", request_id)
if self._storage_source==1: if self._storage_source == 1:
or_future = t.submit(minioSdk.put_object, or_image, or_image_name) or_future = t.submit(minioSdk.put_object, or_image, or_image_name)
else: else:
or_future = t.submit(aliyunOssSdk.put_object, or_image_name, or_image.tobytes()) or_future = t.submit(aliyunOssSdk.put_object, or_image_name, or_image.tobytes())
task.append(or_future) task.append(or_future)
model_info_list = image_result["model_info"] model_info_list = image_result["model_info"]
llm_flag_list = []
msg_list = [] msg_list = []
for model_info in model_info_list: for model_info in model_info_list:
ai_image = cv2.imencode(".jpg", model_info["aFrame"])[1]
fire_flag = model_info.get("modelCode")
print("###line229",fire_flag)
if fire_flag == "008":
aFrame = model_info["aFrame"]
# aFrame_rgb = cv2.cvtColor(aFrame, cv2.COLOR_BGR2RGB)
tar_cls = model_info["detectTargetCode"]
logger.info("目标类别：{}", tar_cls)
# 0715 ： 这块也不对，因为需要具体到目标才可以，否则只看0 1 ，是无法做到判断正确的火情处理，还有大量其他的部分
# is_use_llm = True
if tar_cls == '0':
# tar = "明火"
tar = "2. 判断图中是否有明火。3. 如果有，识别火势的大小，比如：[无情况/小/中/大]"
elif tar_cls == '1':
# tar = "某种物质燃烧产生的烟雾"
tar = "2. 判断图中是否有烟雾。3. 如果有，识别烟雾的颜色，比如：[无烟雾/白/灰白/灰黑]"
else:
self.is_use_llm = False
llm_flag = True
if self.is_use_llm:
# prompt = f"我看到图中有{tar}，请你判断是否正确。简要回答。"
prompt = f"分点进行输出：1. 简短描述图像中的场景。{tar}"
logger.info("检测到目标，大模型提示词：{}", prompt)
# llm_flag, llm_res = self.get_llm_res(aFrame, prompt)
llm_flag, llm_res = self.get_llm_res(aFrame, prompt,self.server_ip,self.server_port)
llm_flag_list.append(llm_flag)
logger.info("检测到目标，大模型识别图中是否存在目标：{}, 大模型输出：{}", llm_flag, llm_res)
if llm_flag:
logger.info("经大模型筛查，小模型识别正确!!!")
# cv2.imwrite(f"/home/thsw2/wei/image_res/{self.cnt}_{tar_cls}_1.jpg", aFrame, [cv2.IMWRITE_JPEG_QUALITY, 80])
# self.cnt = self.cnt + 1
else:
logger.info("经大模型筛查，小模型识别错误!!!")
# cv2.imwrite(f"/home/thsw2/wei/image_res/{self.cnt}_{tar_cls}_0.jpg", aFrame, [cv2.IMWRITE_JPEG_QUALITY, 80])
# self.cnt = self.cnt + 1
else:
llm_flag_list.append(llm_flag)
ai_image = cv2.imencode(".jpg", aFrame)[1]
ai_image_name = build_image_name(image_result["current_frame"], ai_image_name = build_image_name(image_result["current_frame"],
image_result["last_frame"], image_result["last_frame"],
analyse_type, analyse_type,
"AI", "AI",
model_info["modelCode"], model_info["modelCode"],
model_info["detectTargetCode"], model_info["detectTargetCode"],
request_id) request_id)
if llm_flag: # 根据大模型结果判定是否上传 if self._storage_source == 1:
if self._storage_source==1: ai_future = t.submit(minioSdk.put_object, ai_image,
ai_future = t.submit(minioSdk.put_object, ai_image, ai_image_name)
ai_image_name) else:
else: ai_future = t.submit(aliyunOssSdk.put_object, ai_image_name,
ai_future = t.submit(aliyunOssSdk.put_object, ai_image_name, ai_image.tobytes())
ai_image.tobytes())
task.append(ai_future)
# print("###line288",len(task))
#msg_list.append(message_feedback(request_id, task.append(ai_future)
# msg_list.append(message_feedback(request_id,
# AnalysisStatus.RUNNING.value, # AnalysisStatus.RUNNING.value,
# analyse_type, "", "", "", # analyse_type, "", "", "",
# or_image_name, # or_image_name,
# ai_image_name, # ai_image_name,
# model_info['modelCode'], # model_info['modelCode'],
# model_info['detectTargetCode'])) # model_info['detectTargetCode']))
remote_image_list = [] remote_image_list = []
for tk in task: for tk in task:
remote_image_list.append( tk.result() ) remote_image_list.append(tk.result())
for ii, model_info in enumerate(model_info_list): for ii, model_info in enumerate(model_info_list):
msg_list.append(message_feedback(request_id,
if llm_flag_list[ii]: AnalysisStatus.RUNNING.value,
msg_list.append( message_feedback(request_id, analyse_type, "", "", "",
AnalysisStatus.RUNNING.value, remote_image_list[0],
analyse_type, "", "", "", remote_image_list[ii + 1],
remote_image_list[0], model_info['modelCode'],
remote_image_list[ii+1], model_info['detectTargetCode'],
model_info['modelCode'], longitude=model_info['gps'][0],
model_info['detectTargetCode'], latitude=model_info['gps'][1],
longitude=model_info['gps'][0], ))
latitude=model_info['gps'][1],
) ) if (not self._algSwitch) or (self._algStatus and self._algSwitch):
if (not self._algSwitch) or ( self._algStatus and self._algSwitch):
for msg in msg_list: for msg in msg_list:
put_queue(fb_queue, msg, timeout=2, is_ex=False) put_queue(fb_queue, msg, timeout=2, is_ex=False)
del task, msg_list del task, msg_list
@ -348,9 +230,9 @@ def build_image_name(*args):
time_now = TimeUtils.now_date_to_str("%Y-%m-%d-%H-%M-%S") time_now = TimeUtils.now_date_to_str("%Y-%m-%d-%H-%M-%S")
return "%s/%s_frame-%s-%s_type_%s-%s-%s-%s_%s.jpg" % (request_id, time_now, current_frame, last_frame, return "%s/%s_frame-%s-%s_type_%s-%s-%s-%s_%s.jpg" % (request_id, time_now, current_frame, last_frame,
random_num, mode_type, modeCode, target, image_type) random_num, mode_type, modeCode, target, image_type)
#如果任务是图像处理，则用此类
# 如果任务是图像处理，则用此类
class ImageTypeImageFileUpload(Thread): class ImageTypeImageFileUpload(Thread):
__slots__ = ('_fb_queue', '_context', '_image_queue', '_analyse_type', '_msg') __slots__ = ('_fb_queue', '_context', '_image_queue', '_analyse_type', '_msg')
@ -358,6 +240,7 @@ class ImageTypeImageFileUpload(Thread):
super().__init__() super().__init__()
self._fb_queue, self._context, self._msg, self._image_queue, self._analyse_type = args self._fb_queue, self._context, self._msg, self._image_queue, self._analyse_type = args
self._storage_source = self._context['service']['storage_source'] self._storage_source = self._context['service']['storage_source']
@staticmethod @staticmethod
def handle_image(det_xywh, copy_frame, font_config): def handle_image(det_xywh, copy_frame, font_config):
""" """
@ -377,12 +260,21 @@ class ImageTypeImageFileUpload(Thread):
if target_list is not None and len(target_list) > 0: if target_list is not None and len(target_list) > 0:
aiFrame = copy_frame.copy() aiFrame = copy_frame.copy()
for target in target_list: for target in target_list:
draw_painting_joint(target[1], aiFrame, target[3], target[2], target[4], font_config) # 自研车牌模型判断
model_info.append({ if ModelType.CITY_CARPLATE_MODEL.value[1] == str(code):
"modelCode": str(code), draw_name_ocr(target, aiFrame, font_config[cls])
"detectTargetCode": str(cls), elif ModelType.CITY_DENSECROWDCOUNT_MODEL.value[1] == str(code) or \
"frame": aiFrame ModelType.CITY_UNDERBUILDCOUNT_MODEL.value[1] == str(code):
}) draw_name_crowd(target, aiFrame, font_config[cls])
else:
draw_painting_joint(target[1], aiFrame, target[3], target[2], target[4], font_config)
model_info.append({
"modelCode": str(code),
"detectTargetCode": str(cls),
"frame": aiFrame
})
if len(model_info) > 0: if len(model_info) > 0:
image_result = { image_result = {
"or_frame": copy_frame, "or_frame": copy_frame,
@ -402,11 +294,11 @@ class ImageTypeImageFileUpload(Thread):
with ThreadPoolExecutor(max_workers=2) as t: with ThreadPoolExecutor(max_workers=2) as t:
try: try:
# 初始化oss客户端 # 初始化oss客户端
if self._storage_source==1: if self._storage_source == 1:
minioSdk = MinioSdk(base_dir, env, request_id ) minioSdk = MinioSdk(base_dir, env, request_id)
else: else:
aliyunOssSdk = AliyunOssSdk(base_dir, env, request_id) aliyunOssSdk = AliyunOssSdk(base_dir, env, request_id)
start_time = time() start_time = time()
while True: while True:
try: try:
@ -427,15 +319,15 @@ class ImageTypeImageFileUpload(Thread):
if det_xywh is None: if det_xywh is None:
ai_image_name = build_image_name(0, 0, analyse_type, "AI", result.get("modelCode"), ai_image_name = build_image_name(0, 0, analyse_type, "AI", result.get("modelCode"),
result.get("type"), request_id) result.get("type"), request_id)
if self._storage_source==1: if self._storage_source == 1:
ai_future = t.submit(minioSdk.put_object, copy_frame,ai_image_name) ai_future = t.submit(minioSdk.put_object, copy_frame, ai_image_name)
else: else:
ai_future = t.submit(aliyunOssSdk.put_object, ai_image_name, copy_frame) ai_future = t.submit(aliyunOssSdk.put_object, ai_image_name, copy_frame)
task.append(ai_future) task.append(ai_future)
remote_names.append(ai_image_name) remote_names.append(ai_image_name)
#msg_list.append(message_feedback(request_id, # msg_list.append(message_feedback(request_id,
# AnalysisStatus.RUNNING.value, # AnalysisStatus.RUNNING.value,
# analyse_type, "", "", "", # analyse_type, "", "", "",
# image_url, # image_url,
@ -446,17 +338,17 @@ class ImageTypeImageFileUpload(Thread):
else: else:
image_result = self.handle_image(det_xywh, copy_frame, font_config) image_result = self.handle_image(det_xywh, copy_frame, font_config)
if image_result: if image_result:
# 图片帧数编码 # 图片帧数编码
if image_url is None: if image_url is None:
or_result, or_image = cv2.imencode(".jpg", image_result.get("or_frame")) or_result, or_image = cv2.imencode(".jpg", image_result.get("or_frame"))
image_url_0 = build_image_name(image_result.get("current_frame"), image_url_0 = build_image_name(image_result.get("current_frame"),
image_result.get("last_frame"), image_result.get("last_frame"),
analyse_type, analyse_type,
"OR", "0", "O", request_id) "OR", "0", "O", request_id)
if self._storage_source==1: if self._storage_source == 1:
or_future = t.submit(minioSdk.put_object, or_image,image_url_0) or_future = t.submit(minioSdk.put_object, or_image, image_url_0)
else: else:
or_future = t.submit(aliyunOssSdk.put_object, image_url_0, or_future = t.submit(aliyunOssSdk.put_object, image_url_0,
or_image.tobytes()) or_image.tobytes())
@ -472,14 +364,14 @@ class ImageTypeImageFileUpload(Thread):
model_info.get("modelCode"), model_info.get("modelCode"),
model_info.get("detectTargetCode"), model_info.get("detectTargetCode"),
request_id) request_id)
if self._storage_source==1: if self._storage_source == 1:
ai_future = t.submit(minioSdk.put_object, ai_image, ai_image_name) ai_future = t.submit(minioSdk.put_object, ai_image, ai_image_name)
else: else:
ai_future = t.submit(aliyunOssSdk.put_object, ai_image_name, ai_future = t.submit(aliyunOssSdk.put_object, ai_image_name,
ai_image.tobytes()) ai_image.tobytes())
task.append(ai_future) task.append(ai_future)
remote_names.append(ai_image_name) remote_names.append(ai_image_name)
#msg_list.append(message_feedback(request_id, # msg_list.append(message_feedback(request_id,
# AnalysisStatus.RUNNING.value, # AnalysisStatus.RUNNING.value,
# analyse_type, "", "", "", # analyse_type, "", "", "",
# image_url, # image_url,
@ -490,9 +382,8 @@ class ImageTypeImageFileUpload(Thread):
remote_url_list = [] remote_url_list = []
for thread_result in task: for thread_result in task:
remote_url_list.append(thread_result.result()) remote_url_list.append(thread_result.result())
# 以下代码是为了获取图像上传后，返回的全路径地址
#以下代码是为了获取图像上传后，返回的全路径地址
if det_xywh is None: if det_xywh is None:
msg_list.append(message_feedback(request_id, msg_list.append(message_feedback(request_id,
AnalysisStatus.RUNNING.value, AnalysisStatus.RUNNING.value,
@ -505,12 +396,12 @@ class ImageTypeImageFileUpload(Thread):
else: else:
if image_result: if image_result:
if image_url is None: if image_url is None:
for ii in range(len(remote_names)-1): for ii in range(len(remote_names) - 1):
msg_list.append(message_feedback(request_id, msg_list.append(message_feedback(request_id,
AnalysisStatus.RUNNING.value, AnalysisStatus.RUNNING.value,
analyse_type, "", "", "", analyse_type, "", "", "",
remote_url_list[0], remote_url_list[0],
remote_url_list[1+ii], remote_url_list[1 + ii],
model_info.get('modelCode'), model_info.get('modelCode'),
model_info.get('detectTargetCode'), model_info.get('detectTargetCode'),
analyse_results=result)) analyse_results=result))
@ -522,13 +413,10 @@ class ImageTypeImageFileUpload(Thread):
image_url, image_url,
remote_url_list[ii], remote_url_list[ii],
model_info_list[ii].get('modelCode'), model_info_list[ii].get('modelCode'),
model_info_list[ii].get('detectTargetCode'), model_info_list[ii].get(
'detectTargetCode'),
analyse_results=result)) analyse_results=result))
for msg in msg_list: for msg in msg_list:
put_queue(fb_queue, msg, timeout=2, is_ex=False) put_queue(fb_queue, msg, timeout=2, is_ex=False)
else: else:

72
concurrency/IntelligentRecognitionProcess.py
View File

@ -62,8 +62,9 @@ class IntelligentRecognitionProcess(Process):
# 发送waitting消息 # 发送waitting消息
put_queue(self._fb_queue, message_feedback(self._msg["request_id"], AnalysisStatus.WAITING.value, put_queue(self._fb_queue, message_feedback(self._msg["request_id"], AnalysisStatus.WAITING.value,
self._analyse_type, progress=init_progess), timeout=2, is_ex=True) self._analyse_type, progress=init_progess), timeout=2, is_ex=True)
self._storage_source = self._context['service']['storage_source'] self._storage_source = self._context['service']['storage_source']
self._algStatus = False self._algStatus = False
def sendEvent(self, eBody): def sendEvent(self, eBody):
put_queue(self.event_queue, eBody, timeout=2, is_ex=True) put_queue(self.event_queue, eBody, timeout=2, is_ex=True)
@ -92,8 +93,6 @@ class IntelligentRecognitionProcess(Process):
return hb_thread return hb_thread
class OnlineIntelligentRecognitionProcess(IntelligentRecognitionProcess): class OnlineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
__slots__ = () __slots__ = ()
@ -112,19 +111,16 @@ class OnlineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
pullProcess.start() pullProcess.start()
return pullProcess return pullProcess
def upload_video(self,base_dir, env, request_id, orFilePath, aiFilePath): def upload_video(self,base_dir, env, request_id, orFilePath, aiFilePath):
if self._storage_source==1: if self._storage_source==1:
minioSdk = MinioSdk(base_dir, env, request_id ) minioSdk = MinioSdk(base_dir, env, request_id)
upload_video_thread_or = Common(minioSdk.put_object, orFilePath, "or_online_%s.mp4" % request_id) upload_video_thread_or = Common(minioSdk.put_object, orFilePath, "or_online_%s.mp4" % request_id)
upload_video_thread_ai = Common(minioSdk.put_object, aiFilePath, "ai_online_%s.mp4" % request_id) upload_video_thread_ai = Common(minioSdk.put_object, aiFilePath, "ai_online_%s.mp4" % request_id)
else: else:
aliyunVodSdk = ThAliyunVodSdk(base_dir, env, request_id) aliyunVodSdk = ThAliyunVodSdk(base_dir, env, request_id)
upload_video_thread_or = Common(aliyunVodSdk.get_play_url, orFilePath, "or_online_%s" % request_id) upload_video_thread_or = Common(aliyunVodSdk.get_play_url, orFilePath, "or_online_%s" % request_id)
upload_video_thread_ai = Common(aliyunVodSdk.get_play_url, aiFilePath, "ai_online_%s" % request_id) upload_video_thread_ai = Common(aliyunVodSdk.get_play_url, aiFilePath, "ai_online_%s" % request_id)
upload_video_thread_or.setDaemon(True) upload_video_thread_or.setDaemon(True)
upload_video_thread_ai.setDaemon(True) upload_video_thread_ai.setDaemon(True)
upload_video_thread_or.start() upload_video_thread_or.start()
@ -132,6 +128,7 @@ class OnlineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
or_url = upload_video_thread_or.get_result() or_url = upload_video_thread_or.get_result()
ai_url = upload_video_thread_ai.get_result() ai_url = upload_video_thread_ai.get_result()
return or_url, ai_url return or_url, ai_url
''' '''
@staticmethod @staticmethod
def upload_video(base_dir, env, request_id, orFilePath, aiFilePath): def upload_video(base_dir, env, request_id, orFilePath, aiFilePath):
@ -145,7 +142,7 @@ class OnlineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
or_url = upload_video_thread_or.get_result() or_url = upload_video_thread_or.get_result()
ai_url = upload_video_thread_ai.get_result() ai_url = upload_video_thread_ai.get_result()
return or_url, ai_url return or_url, ai_url
''' '''
@staticmethod @staticmethod
def ai_normal_dtection(model, frame, request_id): def ai_normal_dtection(model, frame, request_id):
@ -225,7 +222,7 @@ class OnlineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
ex = None ex = None
# 拉流进程、推流进程、心跳线程 # 拉流进程、推流进程、心跳线程
pull_process, push_process, hb_thread = None, None, None pull_process, push_process, hb_thread = None, None, None
# 事件队列、拉流队列、心跳队列、反馈队列 # 事件队列、拉流队列、心跳队列、反馈队列
event_queue, pull_queue, hb_queue, fb_queue = self.event_queue, self._pull_queue, self._hb_queue, self._fb_queue event_queue, pull_queue, hb_queue, fb_queue = self.event_queue, self._pull_queue, self._hb_queue, self._fb_queue
@ -240,14 +237,14 @@ class OnlineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
# 启动拉流进程(包含拉流线程, 图片上传线程,mqtt读取线程) # 启动拉流进程(包含拉流线程, 图片上传线程,mqtt读取线程)
# 拉流进程初始化时间长, 先启动 # 拉流进程初始化时间长, 先启动
pull_process = self.start_pull_stream(msg, context, fb_queue, pull_queue, image_queue, analyse_type, 25) pull_process = self.start_pull_stream(msg, context, fb_queue, pull_queue, image_queue, analyse_type, 25)
#print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno) #7.0, #print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno) #7.0,
# 启动心跳线程 # 启动心跳线程
hb_thread = self.start_heartbeat(fb_queue, hb_queue, request_id, analyse_type, context) hb_thread = self.start_heartbeat(fb_queue, hb_queue, request_id, analyse_type, context)
# print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno) #7.0, # print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno) #7.0,
# 加载算法模型 # 加载算法模型
model_array = get_model(msg, context, analyse_type) model_array = get_model(msg, context, analyse_type)
#print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno) #9.5 # print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno) #9.5
# 启动推流进程 # 启动推流进程
push_process = self.start_push_stream(msg, push_queue, image_queue, push_ex_queue, hb_queue, context) push_process = self.start_push_stream(msg, push_queue, image_queue, push_ex_queue, hb_queue, context)
#print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno) #print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno)
@ -300,7 +297,8 @@ class OnlineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
for i, model in enumerate(model_array): for i, model in enumerate(model_array):
model_conf, code = model model_conf, code = model
if ModelType.CITY_CARPLATE_MODEL.value[1] == str(code) or \ if ModelType.CITY_CARPLATE_MODEL.value[1] == str(code) or \
ModelType.CITY_DENSECROWDCOUNT_MODEL.value[1] == str(code): ModelType.CITY_DENSECROWDCOUNT_MODEL.value[1] == str(code) or\
ModelType.CITY_UNDERBUILDCOUNT_MODEL.value[1] == str(code):
if draw_config.get(code) is None: if draw_config.get(code) is None:
draw_config[code] = {} draw_config[code] = {}
draw_config["font_config"] = model_conf[4] draw_config["font_config"] = model_conf[4]
@ -310,7 +308,7 @@ class OnlineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
else: else:
model_param = model_conf[1] model_param = model_conf[1]
# (modeType, model_param, allowedList, names, rainbows) # (modeType, model_param, allowedList, names, rainbows)
MODEL_CONFIG[code][2](frame_list[0].shape[1], frame_list[0].shape[0], MODEL_CONFIG[code][2](frame_list[0][0].shape[1], frame_list[0][0].shape[0],
model_conf) model_conf)
if draw_config.get("font_config") is None: if draw_config.get("font_config") is None:
draw_config["font_config"] = model_param['font_config'] draw_config["font_config"] = model_param['font_config']
@ -326,16 +324,16 @@ class OnlineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
# print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno) # print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno)
# 多线程并发处理, 经过测试两个线程最优 # 多线程并发处理, 经过测试两个线程最优
det_array = [] det_array = []
for i, frame in enumerate(frame_list): for i, [frame,_] in enumerate(frame_list):
det_result = t.submit(self.obj_det, self, model_array, frame, task_status, det_result = t.submit(self.obj_det, self, model_array, frame, task_status,
frame_index_list[i], tt, request_id) frame_index_list[i], tt, request_id)
det_array.append(det_result) det_array.append(det_result)
push_objs = [det.result() for det in det_array] push_objs = [det.result() for det in det_array]
#print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno) # print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno)
put_queue(push_queue, put_queue(push_queue,
(1, (frame_list, frame_index_list, all_frames, draw_config, push_objs)), (1, (frame_list, frame_index_list, all_frames, draw_config, push_objs)),
timeout=2, is_ex=True) timeout=2, is_ex=True)
#print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno) # print_cpu_status(requestId=request_id,lineNum=inspect.currentframe().f_lineno)
del det_array, push_objs del det_array, push_objs
del frame_list, frame_index_list, all_frames del frame_list, frame_index_list, all_frames
elif pull_result[0] == 1: elif pull_result[0] == 1:
@ -447,12 +445,12 @@ class OnlineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
class OfflineIntelligentRecognitionProcess(IntelligentRecognitionProcess): class OfflineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
__slots__ = () __slots__ = ()
def upload_video(self,base_dir, env, request_id, aiFilePath): def upload_video(self, base_dir, env, request_id, aiFilePath):
aliyunVodSdk = ThAliyunVodSdk(base_dir, env, request_id) aliyunVodSdk = ThAliyunVodSdk(base_dir, env, request_id)
upload_video_thread_ai = Common(aliyunVodSdk.get_play_url, aiFilePath, "ai_online_%s" % request_id) upload_video_thread_ai = Common(aliyunVodSdk.get_play_url, aiFilePath, "ai_online_%s" % request_id)
if self._storage_source==1: if self._storage_source == 1:
minioSdk = MinioSdk(base_dir, env, request_id ) minioSdk = MinioSdk(base_dir, env, request_id)
upload_video_thread_ai = Common(minioSdk.put_object, aiFilePath, "ai_online_%s.mp4" % request_id) upload_video_thread_ai = Common(minioSdk.put_object, aiFilePath, "ai_online_%s.mp4" % request_id)
else: else:
aliyunVodSdk = ThAliyunVodSdk(base_dir, env, request_id) aliyunVodSdk = ThAliyunVodSdk(base_dir, env, request_id)
@ -473,6 +471,7 @@ class OfflineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
ai_url = upload_video_thread_ai.get_result() ai_url = upload_video_thread_ai.get_result()
return ai_url return ai_url
''' '''
@staticmethod @staticmethod
def ai_normal_dtection(model, frame, request_id): def ai_normal_dtection(model, frame, request_id):
model_conf, code = model model_conf, code = model
@ -626,7 +625,8 @@ class OfflineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
for i, model in enumerate(model_array): for i, model in enumerate(model_array):
model_conf, code = model model_conf, code = model
if ModelType.CITY_CARPLATE_MODEL.value[1] == str(code) or \ if ModelType.CITY_CARPLATE_MODEL.value[1] == str(code) or \
ModelType.CITY_DENSECROWDCOUNT_MODEL.value[1] == str(code): ModelType.CITY_DENSECROWDCOUNT_MODEL.value[1] == str(code) or\
ModelType.CITY_UNDERBUILDCOUNT_MODEL.value[1] == str(code):
if draw_config.get(code) is None: if draw_config.get(code) is None:
draw_config[code] = {} draw_config[code] = {}
draw_config["font_config"] = model_conf[4] draw_config["font_config"] = model_conf[4]
@ -637,7 +637,7 @@ class OfflineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
else: else:
model_param = model_conf[1] model_param = model_conf[1]
# (modeType, model_param, allowedList, names, rainbows) # (modeType, model_param, allowedList, names, rainbows)
MODEL_CONFIG[code][2](frame_list[0].shape[1], frame_list[0].shape[0], MODEL_CONFIG[code][2](frame_list[0][0].shape[1], frame_list[0][0].shape[0],
model_conf) model_conf)
if draw_config.get("font_config") is None: if draw_config.get("font_config") is None:
draw_config["font_config"] = model_param['font_config'] draw_config["font_config"] = model_param['font_config']
@ -651,7 +651,7 @@ class OfflineIntelligentRecognitionProcess(IntelligentRecognitionProcess):
det_array = [] det_array = []
for i, frame in enumerate(frame_list): for i, [frame,_] in enumerate(frame_list):
det_result = t.submit(self.obj_det, self, model_array, frame, task_status, det_result = t.submit(self.obj_det, self, model_array, frame, task_status,
frame_index_list[i], tt, request_id) frame_index_list[i], tt, request_id)
det_array.append(det_result) det_array.append(det_result)
@ -766,7 +766,7 @@ class PhotosIntelligentRecognitionProcess(Process):
put_queue(self._fb_queue, message_feedback(self._msg["request_id"], AnalysisStatus.WAITING.value, put_queue(self._fb_queue, message_feedback(self._msg["request_id"], AnalysisStatus.WAITING.value,
self._analyse_type, progress=init_progess), timeout=2, is_ex=True) self._analyse_type, progress=init_progess), timeout=2, is_ex=True)
self.build_logo(self._msg, self._context) self.build_logo(self._msg, self._context)
self._storage_source = self._context['service']['storage_source'] self._storage_source = self._context['service']['storage_source']
@staticmethod @staticmethod
def build_logo(msg, context): def build_logo(msg, context):
@ -943,7 +943,7 @@ class PhotosIntelligentRecognitionProcess(Process):
logger.error("模型分析异常: {}, requestId: {}", format_exc(), request_id) logger.error("模型分析异常: {}, requestId: {}", format_exc(), request_id)
raise e raise e
#自研究车牌模型 # 自研究车牌模型
def carplate_rec(self, imageUrl, mod, image_queue, request_id): def carplate_rec(self, imageUrl, mod, image_queue, request_id):
try: try:
# model_conf： modeType, allowedList, detpar, ocrmodel, rainbows # model_conf： modeType, allowedList, detpar, ocrmodel, rainbows
@ -982,7 +982,6 @@ class PhotosIntelligentRecognitionProcess(Process):
# param = [image, new_device, model, par, img_type, request_id] # param = [image, new_device, model, par, img_type, request_id]
# model_conf, frame, device, requestId # model_conf, frame, device, requestId
dataBack = MODEL_CONFIG[code][3]([[modeType, device, model, postPar], image, request_id])[0][2] dataBack = MODEL_CONFIG[code][3]([[modeType, device, model, postPar], image, request_id])[0][2]
logger.info("当前人数：{}", dataBack[0][0])
dets[code][0] = dataBack dets[code][0] = dataBack
if not dataBack: if not dataBack:
logger.info("当前页面无人") logger.info("当前页面无人")
@ -1055,6 +1054,8 @@ class PhotosIntelligentRecognitionProcess(Process):
ai_result_list = p_result[2] ai_result_list = p_result[2]
for ai_result in ai_result_list: for ai_result in ai_result_list:
box, score, cls = xywh2xyxy2(ai_result) box, score, cls = xywh2xyxy2(ai_result)
if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code):
box.append(ai_result[-1])
# 如果检测目标在识别任务中,继续处理 # 如果检测目标在识别任务中,继续处理
if cls in allowedList: if cls in allowedList:
label_array = label_arraylist[cls] label_array = label_arraylist[cls]
@ -1212,7 +1213,8 @@ class PhotosIntelligentRecognitionProcess(Process):
image_thread.setDaemon(True) image_thread.setDaemon(True)
image_thread.start() image_thread.start()
return image_thread return image_thread
def check_ImageUrl_Vaild(self,url,timeout=1):
def check_ImageUrl_Vaild(self, url, timeout=1):
try: try:
# 发送 HTTP 请求，尝试访问图片 # 发送 HTTP 请求，尝试访问图片
response = requests.get(url, timeout=timeout) # 设置超时时间为 10 秒 response = requests.get(url, timeout=timeout) # 设置超时时间为 10 秒
@ -1243,8 +1245,7 @@ class PhotosIntelligentRecognitionProcess(Process):
ExceptionType.URL_ADDRESS_ACCESS_FAILED.value[0], ExceptionType.URL_ADDRESS_ACCESS_FAILED.value[0],
ExceptionType.URL_ADDRESS_ACCESS_FAILED.value[1]), timeout=2) ExceptionType.URL_ADDRESS_ACCESS_FAILED.value[1]), timeout=2)
return return
with ThreadPoolExecutor(max_workers=1) as t: with ThreadPoolExecutor(max_workers=1) as t:
try: try:
@ -1271,7 +1272,8 @@ class PhotosIntelligentRecognitionProcess(Process):
result = t.submit(self.carpalteRec, imageUrls, model, image_queue, request_id) result = t.submit(self.carpalteRec, imageUrls, model, image_queue, request_id)
task_list.append(result) task_list.append(result)
# 人群计数模型 # 人群计数模型
elif model[1] == ModelType.CITY_DENSECROWDCOUNT_MODEL.value[1]: elif model[1] == ModelType.CITY_DENSECROWDCOUNT_MODEL.value[1] or \
model[1] == ModelType.CITY_UNDERBUILDCOUNT_MODEL.value[1]:
result = t.submit(self.denscrowdcountRec, imageUrls, model, image_queue, request_id) result = t.submit(self.denscrowdcountRec, imageUrls, model, image_queue, request_id)
task_list.append(result) task_list.append(result)
else: else:
@ -1321,6 +1323,7 @@ class ScreenRecordingProcess(Process):
recording_feedback(self._msg["request_id"], RecordingStatus.RECORDING_WAITING.value[0]), recording_feedback(self._msg["request_id"], RecordingStatus.RECORDING_WAITING.value[0]),
timeout=1, is_ex=True) timeout=1, is_ex=True)
self._storage_source = self._context['service']['storage_source'] self._storage_source = self._context['service']['storage_source']
def sendEvent(self, result): def sendEvent(self, result):
put_queue(self._event_queue, result, timeout=2, is_ex=True) put_queue(self._event_queue, result, timeout=2, is_ex=True)
@ -1486,9 +1489,9 @@ class ScreenRecordingProcess(Process):
clear_queue(self._hb_queue) clear_queue(self._hb_queue)
clear_queue(self._pull_queue) clear_queue(self._pull_queue)
def upload_video(self,base_dir, env, request_id, orFilePath): def upload_video(self, base_dir, env, request_id, orFilePath):
if self._storage_source==1: if self._storage_source == 1:
minioSdk = MinioSdk(base_dir, env, request_id ) minioSdk = MinioSdk(base_dir, env, request_id)
upload_video_thread_ai = Common(minioSdk.put_object, aiFilePath, "%s/ai_online.mp4" % request_id) upload_video_thread_ai = Common(minioSdk.put_object, aiFilePath, "%s/ai_online.mp4" % request_id)
else: else:
aliyunVodSdk = ThAliyunVodSdk(base_dir, env, request_id) aliyunVodSdk = ThAliyunVodSdk(base_dir, env, request_id)
@ -1498,6 +1501,7 @@ class ScreenRecordingProcess(Process):
upload_video_thread_ai.start() upload_video_thread_ai.start()
or_url = upload_video_thread_ai.get_result() or_url = upload_video_thread_ai.get_result()
return or_url return or_url
''' '''
@staticmethod @staticmethod
def upload_video(base_dir, env, request_id, orFilePath): def upload_video(base_dir, env, request_id, orFilePath):

303
concurrency/PullMqttThread.py
View File

@ -1,142 +1,163 @@
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
from threading import Thread from threading import Thread
from time import sleep, time from time import sleep, time
from traceback import format_exc from traceback import format_exc
from loguru import logger from loguru import logger
from common.YmlConstant import mqtt_yml_path from common.YmlConstant import mqtt_yml_path
from util.RWUtils import getConfigs from util.RWUtils import getConfigs
from common.Constant import init_progess from common.Constant import init_progess
from enums.AnalysisStatusEnum import AnalysisStatus from enums.AnalysisStatusEnum import AnalysisStatus
from entity.FeedBack import message_feedback from entity.FeedBack import message_feedback
from enums.ExceptionEnum import ExceptionType from enums.ExceptionEnum import ExceptionType
from exception.CustomerException import ServiceException from exception.CustomerException import ServiceException
from util.QueUtil import get_no_block_queue, put_queue, clear_queue from util.QueUtil import get_no_block_queue, put_queue, clear_queue
from multiprocessing import Process, Queue from multiprocessing import Process, Queue
import paho.mqtt.client as mqtt import paho.mqtt.client as mqtt
import json,os import json,os
class PullMqtt(Thread): class PullMqtt(Thread):
__slots__ = ('__fb_queue', '__mqtt_list', '__request_id', '__analyse_type', "_context") __slots__ = ('__fb_queue', '__mqtt_list', '__request_id', '__analyse_type', "_context" ,'__business')
def __init__(self, *args): def __init__(self, *args):
super().__init__() super().__init__()
self.__fb_queue, self.__mqtt_list, self.__request_id, self.__analyse_type, self._context = args self.__fb_queue, self.__mqtt_list, self.__request_id, self.__analyse_type, self._context, self.__business = args
base_dir, env = self._context["base_dir"], self._context["env"] base_dir, env = self._context["base_dir"], self._context["env"]
self.__config = getConfigs(os.path.join(base_dir, mqtt_yml_path % env)) self.__config = getConfigs(os.path.join(base_dir, mqtt_yml_path % env))
if self.__business == 0:
self.__broker = self.__config["broker"] self.__broker = self.__config['location']["broker"]
self.__port = self.__config["port"] self.__port = self.__config['location']["port"]
self.__topic = self.__config["topic"] self.__topic = self.__config['location']["topic"]
self.__lengthMqttList = self.__config["length"] elif self.__business == 1:
self.__broker = self.__config['invade']["broker"]
self.__port = self.__config['invade']["port"]
def put_queue(self,__queue,data): self.__topic = self.__config['invade']["topic"]
if __queue.full(): self.__lengthMqttList = self.__config["length"]
a = __queue.get()
__queue.put( data,block=True, timeout=2 )
def on_connect(self,client,userdata,flags,rc): def put_queue(self,__queue,data):
client.subscribe(self.__topic) if __queue.full():
a = __queue.get()
__queue.put( data,block=True, timeout=2 )
def on_connect(self,client,userdata,flags,rc):
# 当接收到MQTT消息时，回调函数 client.subscribe(self.__topic)
def on_message(self,client, userdata, msg):
# 将消息解码为JSON格式
payload = msg.payload.decode('utf-8')
data = json.loads(payload) # 当接收到MQTT消息时，回调函数
#logger.info(str(data)) def on_location(self,client, userdata, msg):
# 将消息解码为JSON格式
payload = msg.payload.decode('utf-8')
# 解析位姿信息 data = json.loads(payload)
lon = data.get("lon") #logger.info(str(data))
lat = data.get("lat") # 解析位姿信息
alt = data.get("alt") lon = data.get("lon")
yaw = data.get("yaw") lat = data.get("lat")
pitch = data.get("pitch") alt = data.get("alt")
roll = data.get("roll") yaw = data.get("yaw")
pitch = data.get("pitch")
if len(self.__mqtt_list) == self.__lengthMqttList: roll = data.get("roll")
self.__mqtt_list.pop(0)
self.__mqtt_list.append(data) if len(self.__mqtt_list) == self.__lengthMqttList:
self.__mqtt_list.pop(0)
self.__mqtt_list.append([self.__business,data])
# 打印无人机的位姿信息
#print(f"Longitude: {lon}, Latitude: {lat}, Altitude: {alt}, sat:{data.get('satcount')} , list length:{len(self.__mqtt_list)}")
# 打印无人机的位姿信息
def mqtt_connect(self): #print(f"Longitude: {lon}, Latitude: {lat}, Altitude: {alt}, sat:{data.get('satcount')} , list length:{len(self.__mqtt_list)}")
# 创建客户端
self.client = mqtt.Client() def on_invade(self, client, userdata, msg):
self.client.on_connect = self.on_connect # 将消息解码为JSON格式
# 设置回调函数 payload = msg.payload.decode('utf-8')
self.client.on_message = self.on_message data = json.loads(payload)
# logger.info(str(data))
# 连接到 Broker # 解析位姿信息
self.client.connect(self.__broker, self.__port) points = data.get("points")
# 订阅主题 if len(self.__mqtt_list) == self.__lengthMqttList:
self.client.subscribe(self.__topic) self.__mqtt_list.pop(0)
# 循环等待并处理网络事件 self.__mqtt_list.append([self.__business,points])
self.client.loop_forever()
# 打印无人机的位姿信息
def mqtt_disconnect(self): # print(f"Longitude: {lon}, Latitude: {lat}, Altitude: {alt}, sat:{data.get('satcount')} , list length:{len(self.__mqtt_list)}")
start_time = time()
while True: def mqtt_connect(self):
if time() - start_time > service_timeout: # 创建客户端
logger.error("MQTT读取超时, requestId: %s,限定时间：%.1s , 已运行：%.1fs"%(request_id,service_timeout, time() - start_time)) self.client = mqtt.Client()
raise ServiceException(ExceptionType.TASK_EXCUTE_TIMEOUT.value[0], self.client.on_connect = self.on_connect
ExceptionType.TASK_EXCUTE_TIMEOUT.value[1]) if self.__business == 0:
client.loop_stop() # 停止循环 # 设置回调函数
client.disconnect() # 断开连接 self.client.on_message = self.on_location
elif self.__business == 1:
def run(self): # 设置回调函数
request_id, mqtt_list, progress = self.__request_id, self.__mqtt_list, init_progess self.client.on_message = self.on_invade
analyse_type, fb_queue = self.__analyse_type, self.__fb_queue
#service_timeout = int(self.__config["service"]["timeout"]) + 120 # 连接到 Broker
self.client.connect(self.__broker, self.__port)
try:
logger.info("开始MQTT读取线程！requestId:{}", request_id) # 订阅主题
mqtt_init_num = 0 self.client.subscribe(self.__topic)
self.mqtt_connect() # 循环等待并处理网络事件
self.client.loop_forever()
except Exception:
logger.error("MQTT线程异常:{}, requestId:{}", format_exc(), request_id) def mqtt_disconnect(self):
finally: start_time = time()
mqtt_list = [] while True:
logger.info("MQTT线程停止完成！requestId:{}", request_id) if time() - start_time > service_timeout:
logger.error("MQTT读取超时, requestId: %s,限定时间：%.1s , 已运行：%.1fs"%(request_id,service_timeout, time() - start_time))
raise ServiceException(ExceptionType.TASK_EXCUTE_TIMEOUT.value[0],
def start_PullMqtt(fb_queue, mqtt_list, request_id, analyse_type, context): ExceptionType.TASK_EXCUTE_TIMEOUT.value[1])
mqtt_thread = PullMqtt(fb_queue, mqtt_list, request_id, analyse_type, context) client.loop_stop() # 停止循环
mqtt_thread.setDaemon(True) client.disconnect() # 断开连接
mqtt_thread.start()
return mqtt_thread def run(self):
def start_PullVideo(mqtt_list): request_id, mqtt_list, progress = self.__request_id, self.__mqtt_list, init_progess
for i in range(1000): analyse_type, fb_queue = self.__analyse_type, self.__fb_queue
sleep(1) #service_timeout = int(self.__config["service"]["timeout"]) + 120
if len(mqtt_list)>=10:
print( mqtt_list[4]) try:
print(i,len(mqtt_list)) logger.info("开始MQTT读取线程！requestId:{}", request_id)
if __name__=="__main__": mqtt_init_num = 0
#context = {'service':{'timeout':3600},'mqtt':{ self.mqtt_connect()
# 'broker':"101.133.163.127",'port':1883,'topic':"test/topic","length":10}
# } except Exception:
context = { logger.error("MQTT线程异常:{}, requestId:{}", format_exc(), request_id)
'base_dir':'/home/th/WJ/test/tuoheng_algN', finally:
'env':'test' mqtt_list = []
logger.info("MQTT线程停止完成！requestId:{}", request_id)
}
analyse_type = '1'
request_id = '123456789' def start_PullMqtt(fb_queue, mqtt_list, request_id, analyse_type, context):
event_queue, pull_queue, mqtt_list, image_queue, push_queue, push_ex_queue = Queue(), Queue(10), [], Queue(), Queue(), Queue() mqtt_thread = PullMqtt(fb_queue, mqtt_list, request_id, analyse_type, context)
fb_queue = Queue() mqtt_thread.setDaemon(True)
mqtt_thread = start_PullMqtt(fb_queue, mqtt_list, request_id, analyse_type, context) mqtt_thread.start()
return mqtt_thread
def start_PullVideo(mqtt_list):
start_PullVideo(mqtt_list) for i in range(1000):
print('---line117--') sleep(1)
if len(mqtt_list)>=10:
print( mqtt_list[4])
print(i,len(mqtt_list))
#mqtt_thread.join() if __name__=="__main__":
#context = {'service':{'timeout':3600},'mqtt':{
# 'broker':"101.133.163.127",'port':1883,'topic':"test/topic","length":10}
# }
context = {
'base_dir':'/home/th/WJ/test/tuoheng_algN',
'env':'test'
}
analyse_type = '1'
request_id = '123456789'
event_queue, pull_queue, mqtt_list, image_queue, push_queue, push_ex_queue = Queue(), Queue(10), [], Queue(), Queue(), Queue()
fb_queue = Queue()
mqtt_thread = start_PullMqtt(fb_queue, mqtt_list, request_id, analyse_type, context)
start_PullVideo(mqtt_list)
print('---line117--')
#mqtt_thread.join()

40
concurrency/PullVideoStreamProcess.py
View File

@ -35,15 +35,15 @@ class PullVideoStreamProcess(Process):
put_queue(self._command_queue, result, timeout=2, is_ex=True) put_queue(self._command_queue, result, timeout=2, is_ex=True)
@staticmethod @staticmethod
def start_File_upload(fb_queue, context, msg, image_queue, analyse_type,mqtt_list): def start_File_upload(fb_queue, context, msg, image_queue, analyse_type):
image_thread = ImageFileUpload(fb_queue, context, msg, image_queue, analyse_type,mqtt_list) image_thread = ImageFileUpload(fb_queue, context, msg, image_queue, analyse_type)
image_thread.setDaemon(True) image_thread.setDaemon(True)
image_thread.start() image_thread.start()
return image_thread return image_thread
@staticmethod @staticmethod
def start_PullMqtt(fb_queue, mqtt_list, request_id, analyse_type, context): def start_PullMqtt(fb_queue, mqtt_list, request_id, analyse_type, context,business):
mqtt_thread = PullMqtt(fb_queue, mqtt_list, request_id, analyse_type, context) mqtt_thread = PullMqtt(fb_queue, mqtt_list, request_id, analyse_type, context,business)
mqtt_thread.setDaemon(True) mqtt_thread.setDaemon(True)
mqtt_thread.start() mqtt_thread.start()
return mqtt_thread return mqtt_thread
@ -81,13 +81,14 @@ class OnlinePullVideoStreamProcess(PullVideoStreamProcess):
# 初始化日志 # 初始化日志
init_log(base_dir, env) init_log(base_dir, env)
logger.info("开启启动实时视频拉流进程, requestId:{},pid:{},ppid:{}", request_id,os.getpid(),os.getppid()) logger.info("开启启动实时视频拉流进程, requestId:{},pid:{},ppid:{}", request_id,os.getpid(),os.getppid())
#开启mqtt # 开启mqtt
if service["mqtt_flag"]==1: if service['mqtt']["flag"] == 1:
mqtt_thread = self.start_PullMqtt(fb_queue, mqtt_list, request_id, analyse_type, context) business = service['mqtt']["business"]
mqtt_thread = self.start_PullMqtt(fb_queue, mqtt_list, request_id, analyse_type, context, business)
# 开启图片上传线程 # 开启图片上传线程
image_thread = self.start_File_upload(fb_queue, context, msg, image_queue, analyse_type,mqtt_list) image_thread = self.start_File_upload(fb_queue, context, msg, image_queue, analyse_type)
cv2_init_num, init_pull_num, concurrent_frame = 0, 1, 1 cv2_init_num, init_pull_num, concurrent_frame = 0, 1, 1
start_time, pull_stream_start_time, read_start_time, full_timeout = time(), None, None, None start_time, pull_stream_start_time, read_start_time, full_timeout = time(), None, None, None
while True: while True:
@ -129,7 +130,7 @@ class OnlinePullVideoStreamProcess(PullVideoStreamProcess):
frame, pull_p, width, height = pull_read_video_stream(pull_p, pull_url, width, height, width_height_3, frame, pull_p, width, height = pull_read_video_stream(pull_p, pull_url, width, height, width_height_3,
w_2, h_2, request_id) w_2, h_2, request_id)
if pull_queue.full(): if pull_queue.full():
logger.info("pull拉流队列满了:{}, requestId: {}", os.getppid(), request_id) #logger.info("pull拉流队列满了:{}, requestId: {}", os.getppid(), request_id)
if full_timeout is None: if full_timeout is None:
full_timeout = time() full_timeout = time()
if time() - full_timeout > 180: if time() - full_timeout > 180:
@ -171,7 +172,7 @@ class OnlinePullVideoStreamProcess(PullVideoStreamProcess):
sleep(1) sleep(1)
continue continue
init_pull_num, read_start_time = 1, None init_pull_num, read_start_time = 1, None
frame_list.append(frame) frame_list.append([frame, mqtt_list])
frame_index_list.append(concurrent_frame) frame_index_list.append(concurrent_frame)
if len(frame_list) >= frame_num: if len(frame_list) >= frame_num:
put_queue(pull_queue, (4, (frame_list, frame_index_list, all_frames)), timeout=1, is_ex=True) put_queue(pull_queue, (4, (frame_list, frame_index_list, all_frames)), timeout=1, is_ex=True)
@ -222,10 +223,11 @@ class OfflinePullVideoStreamProcess(PullVideoStreamProcess):
def run(self): def run(self):
msg, context, frame_num, analyse_type = self._msg, self._context, self._frame_num, self._analyse_type msg, context, frame_num, analyse_type = self._msg, self._context, self._frame_num, self._analyse_type
request_id, base_dir, env, pull_url = msg["request_id"], context['base_dir'], context['env'], msg["pull_url"] request_id, base_dir, env, pull_url, service = msg["request_id"], context['base_dir'], context['env'], msg["pull_url"], context["service"]
ex, service_timeout, full_timeout = None, int(context["service"]["timeout"]) + 120, None ex, service_timeout, full_timeout = None, int(context["service"]["timeout"]) + 120, None
command_queue, pull_queue, image_queue, fb_queue = self._command_queue, self._pull_queue, self._image_queue, \
self._fb_queue command_queue, pull_queue, image_queue, fb_queue, mqtt_list = self._command_queue, self._pull_queue, self._image_queue, \
self._fb_queue, self._mqtt_list
image_thread, pull_p = None, None image_thread, pull_p = None, None
width, height, width_height_3, all_frames, w_2, h_2 = None, None, None, 0, None, None width, height, width_height_3, all_frames, w_2, h_2 = None, None, None, 0, None, None
frame_list, frame_index_list = [], [] frame_list, frame_index_list = [], []
@ -235,8 +237,12 @@ class OfflinePullVideoStreamProcess(PullVideoStreamProcess):
init_log(base_dir, env) init_log(base_dir, env)
logger.info("开启离线视频拉流进程, requestId:{}", request_id) logger.info("开启离线视频拉流进程, requestId:{}", request_id)
#开启mqtt
if service['mqtt']["flag"]==1:
business = service['mqtt']["business"]
mqtt_thread = self.start_PullMqtt(fb_queue, mqtt_list, request_id, analyse_type, context, business)
# 开启图片上传线程 # 开启图片上传线程
image_thread = self.start_File_upload(fb_queue, context, msg, image_queue, analyse_type,[]) image_thread = self.start_File_upload(fb_queue, context, msg, image_queue, analyse_type)
# 初始化拉流工具类 # 初始化拉流工具类
cv2_init_num, concurrent_frame = 0, 1 cv2_init_num, concurrent_frame = 0, 1
@ -269,7 +275,7 @@ class OfflinePullVideoStreamProcess(PullVideoStreamProcess):
width, height, width_height_3, all_frames, w_2, h_2 = build_video_info(pull_url, request_id) width, height, width_height_3, all_frames, w_2, h_2 = build_video_info(pull_url, request_id)
continue continue
if pull_queue.full(): if pull_queue.full():
logger.info("pull拉流队列满了:{}, requestId: {}", os.getppid(), request_id) #logger.info("pull拉流队列满了:{}, requestId: {}", os.getppid(), request_id)
if full_timeout is None: if full_timeout is None:
full_timeout = time() full_timeout = time()
if time() - full_timeout > 180: if time() - full_timeout > 180:
@ -306,7 +312,7 @@ class OfflinePullVideoStreamProcess(PullVideoStreamProcess):
ExceptionType.READSTREAM_TIMEOUT_EXCEPTION.value[1]) ExceptionType.READSTREAM_TIMEOUT_EXCEPTION.value[1])
logger.info("离线拉流线程结束, requestId: {}", request_id) logger.info("离线拉流线程结束, requestId: {}", request_id)
break break
frame_list.append(frame) frame_list.append([frame,mqtt_list])
frame_index_list.append(concurrent_frame) frame_index_list.append(concurrent_frame)
if len(frame_list) >= frame_num: if len(frame_list) >= frame_num:
put_queue(pull_queue, (4, (frame_list, frame_index_list, all_frames)), timeout=1, is_ex=True) put_queue(pull_queue, (4, (frame_list, frame_index_list, all_frames)), timeout=1, is_ex=True)

255
concurrency/PushVideoStreamProcess.py
View File

@ -23,7 +23,7 @@ from util.Cv2Utils import video_conjuncing, write_or_video, write_ai_video, push
from util.ImageUtils import url2Array, add_water_pic from util.ImageUtils import url2Array, add_water_pic
from util.LogUtils import init_log from util.LogUtils import init_log
from util.PlotsUtils import draw_painting_joint, filterBox, xywh2xyxy2, draw_name_joint from util.PlotsUtils import draw_painting_joint, filterBox, xywh2xyxy2, xy2xyxy, draw_name_joint, plot_one_box_auto, draw_name_ocr,draw_name_crowd,draw_transparent_red_polygon
from util.QueUtil import get_no_block_queue, put_queue, clear_queue from util.QueUtil import get_no_block_queue, put_queue, clear_queue
@ -36,11 +36,10 @@ class PushStreamProcess(Process):
# 传参 # 传参
self._msg, self._push_queue, self._image_queue, self._push_ex_queue, self._hb_queue, self._context = args self._msg, self._push_queue, self._image_queue, self._push_ex_queue, self._hb_queue, self._context = args
self._algStatus = False # 默认关闭 self._algStatus = False # 默认关闭
self._algSwitch = self._context['service']['algSwitch'] self._algSwitch = self._context['service']['algSwitch']
# 0521:
#0521: default_enabled = str(self._msg.get("defaultEnabled", "True")).lower() == "true"
default_enabled = str(self._msg.get("defaultEnabled", "True")).lower() == "true"
if default_enabled: if default_enabled:
print("执行默认程序（defaultEnabled=True）") print("执行默认程序（defaultEnabled=True）")
self._algSwitch = True self._algSwitch = True
@ -49,16 +48,9 @@ class PushStreamProcess(Process):
print("执行替代程序（defaultEnabled=False）") print("执行替代程序（defaultEnabled=False）")
# 这里放非默认逻辑的代码 # 这里放非默认逻辑的代码
self._algSwitch = False self._algSwitch = False
print("---line53 :PushVideoStreamProcess.py---",self._algSwitch) print("---line53 :PushVideoStreamProcess.py---",self._algSwitch)
def build_logo_url(self): def build_logo_url(self):
logo = None logo = None
if self._context["video"]["video_add_water"]: if self._context["video"]["video_add_water"]:
@ -138,15 +130,26 @@ class OnPushStreamProcess(PushStreamProcess):
if push_r is not None: if push_r is not None:
if push_r[0] == 1: if push_r[0] == 1:
frame_list, frame_index_list, all_frames, draw_config, push_objs = push_r[1] frame_list, frame_index_list, all_frames, draw_config, push_objs = push_r[1]
for i, frame in enumerate(frame_list): # 处理每1帧
for i, [frame,mqtt_list] in enumerate(frame_list):
# mqtt传参
border = None
mqttPares = None
if len(mqtt_list) >= 1:
mqttPares = mqtt_list[0]
if mqttPares[0] == 1:
border = mqttPares[1]
pix_dis = int((frame.shape[0]//10)*1.2) pix_dis = int((frame.shape[0]//10)*1.2)
# 复制帧用来画图 # 复制帧用来画图
copy_frame = frame.copy() copy_frame = frame.copy()
if border is not None:
copy_frame = draw_transparent_red_polygon(copy_frame, np.array(border, np.int32),alpha=0.25)
det_xywh, thread_p = {}, [] det_xywh, thread_p = {}, []
det_xywh2 = {} det_xywh2 = {'mqttPares':mqttPares}
# 所有问题的矩阵集合 # 所有问题的矩阵集合
qs_np = None qs_np = None
qs_reurn = [] qs_reurn = []
bp_np = None
for det in push_objs[i]: for det in push_objs[i]:
code, det_result = det code, det_result = det
# 每个单独模型处理 # 每个单独模型处理
@ -155,17 +158,42 @@ class OnPushStreamProcess(PushStreamProcess):
font_config, allowedList = draw_config["font_config"], draw_config[code]["allowedList"] font_config, allowedList = draw_config["font_config"], draw_config[code]["allowedList"]
rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"] rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"]
for qs in det_result: for qs in det_result:
try: # 应对NaN情况 # 自研车牌模型处理
box, score, cls = xywh2xyxy2(qs) if ModelType.CITY_CARPLATE_MODEL.value[1] == str(code):
except: cls = 0
continue box = xy2xyxy(qs[1])
if cls not in allowedList or score < frame_score: score = None
continue color = rainbows[cls]
label_array, color = label_arrays[cls], rainbows[cls] label_array = None
if ModelType.CHANNEL2_MODEL.value[1] == str(code) and cls == 2: rr = t.submit(draw_name_ocr, qs, copy_frame, color)
rr = t.submit(draw_name_joint, box, copy_frame, draw_config[code]["label_dict"], score, color, font_config, qs[6]) elif ModelType.CITY_DENSECROWDCOUNT_MODEL.value[1] == str(code) or\
ModelType.CITY_UNDERBUILDCOUNT_MODEL.value[1] == str(code):
cls = 0
# crowdlabel, points = qs
box = [(0, 0), (0, 0), (0, 0), (0, 0)]
score = None
color = rainbows[cls]
label_array = None
rr = t.submit(draw_name_crowd, qs, copy_frame, color)
else: else:
rr = t.submit(draw_painting_joint, box, copy_frame, label_array, score, color, font_config) try: # 应对NaN情况
box, score, cls = xywh2xyxy2(qs)
if cls not in allowedList or score < frame_score:
continue
if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code):
# 借score作为points点集
box.append(qs[-1])
except:
continue
label_array, color = label_arrays[cls], rainbows[cls]
if ModelType.CHANNEL2_MODEL.value[1] == str(code) and cls == 2:
rr = t.submit(draw_name_joint, box, copy_frame,
draw_config[code]["label_dict"], score, color,
font_config, qs[6])
else:
rr = t.submit(draw_painting_joint, box, copy_frame, label_array, score, color, font_config, border=border)
thread_p.append(rr) thread_p.append(rr)
if det_xywh.get(code) is None: if det_xywh.get(code) is None:
det_xywh[code] = {} det_xywh[code] = {}
@ -174,32 +202,35 @@ class OnPushStreamProcess(PushStreamProcess):
if cd is None: if cd is None:
det_xywh[code][cls] = [[cls, box, score, label_array, color]] det_xywh[code][cls] = [[cls, box, score, label_array, color]]
else: else:
det_xywh[code][cls].append([cls, box, score, label_array, color]) det_xywh[code][cls].append([cls, box, score, label_array, color])
if qs_np is None: if qs_np is None:
qs_np = np.array([box[0][0], box[0][1], box[1][0], box[1][1], qs_np = np.array([box[0][0], box[0][1], box[1][0], box[1][1],
box[2][0], box[2][1], box[3][0], box[3][1], box[2][0], box[2][1], box[3][0], box[3][1],
score, cls, code],dtype=np.float32) score, cls, code],dtype=np.float32)
else: else:
result_li = np.array([box[0][0], box[0][1], box[1][0], box[1][1], result_li = np.array([box[0][0], box[0][1], box[1][0], box[1][1],
box[2][0], box[2][1], box[3][0], box[3][1], box[2][0], box[2][1], box[3][0], box[3][1],
score, cls, code],dtype=np.float32) score, cls, code],dtype=np.float32)
qs_np = np.row_stack((qs_np, result_li)) qs_np = np.row_stack((qs_np, result_li))
if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code):
if bp_np is None:
bp_np = np.array([box[0][0], box[0][1], box[-1]], dtype=object)
else:
bp_li = np.array([box[0][0], box[0][1], box[-1]], dtype=object)
bp_np = np.row_stack((bp_np, bp_li))
if logo: if logo:
frame = add_water_pic(frame, logo, request_id) frame = add_water_pic(frame, logo, request_id)
copy_frame = add_water_pic(copy_frame, logo, request_id) copy_frame = add_water_pic(copy_frame, logo, request_id)
if len(thread_p) > 0: if len(thread_p) > 0:
for r in thread_p: for r in thread_p:
r.result() r.result()
#print('----line173:',self._algSwitch,self._algStatus) #print('----line173:',self._algSwitch,self._algStatus)
if self._algSwitch and (not self._algStatus): if self._algSwitch and (not self._algStatus):
# frame_merge = video_conjuncing(frame, frame.copy()) frame_merge = video_conjuncing(frame, frame.copy())
else:
frame_merge = frame.copy() frame_merge = video_conjuncing(frame, copy_frame)
else:
# frame_merge = video_conjuncing(frame, copy_frame)
frame_merge = copy_frame
# 写原视频到本地 # 写原视频到本地
write_or_video_result = t.submit(write_or_video, frame, orFilePath, or_video_file, write_or_video_result = t.submit(write_or_video, frame, orFilePath, or_video_file,
or_write_status, request_id) or_write_status, request_id)
@ -211,7 +242,7 @@ class OnPushStreamProcess(PushStreamProcess):
# 如果有问题, 走下面的逻辑 # 如果有问题, 走下面的逻辑
if qs_np is not None: if qs_np is not None:
if len(qs_np.shape) == 1: if len(qs_np.shape) == 1:
qs_np = qs_np[np.newaxis,...] qs_np = qs_np[np.newaxis,...]
qs_np_id = qs_np.copy() qs_np_id = qs_np.copy()
b = np.ones(qs_np_id.shape[0]) b = np.ones(qs_np_id.shape[0])
qs_np_id = np.column_stack((qs_np_id,b)) qs_np_id = np.column_stack((qs_np_id,b))
@ -219,7 +250,7 @@ class OnPushStreamProcess(PushStreamProcess):
if picture_similarity: if picture_similarity:
qs_np_tmp = qs_np_id.copy() qs_np_tmp = qs_np_id.copy()
b = np.zeros(qs_np.shape[0]) b = np.zeros(qs_np.shape[0])
qs_reurn = np.column_stack((qs_np,b)) qs_reurn = np.column_stack((qs_np,b))
else: else:
qs_reurn = filterBox(qs_np, qs_np_tmp, pix_dis) qs_reurn = filterBox(qs_np, qs_np_tmp, pix_dis)
if picture_similarity: if picture_similarity:
@ -237,7 +268,7 @@ class OnPushStreamProcess(PushStreamProcess):
if q[11] >= 1: if q[11] >= 1:
cls = int(q[9]) cls = int(q[9])
if not (cls in new_lab): if not (cls in new_lab):
continue # 为了防止其他类别被带出 continue # 为了防止其他类别被带出
code = str(int(q[10])).zfill(3) code = str(int(q[10])).zfill(3)
if det_xywh2.get(code) is None: if det_xywh2.get(code) is None:
det_xywh2[code] = {} det_xywh2[code] = {}
@ -245,17 +276,29 @@ class OnPushStreamProcess(PushStreamProcess):
score = q[8] score = q[8]
rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"] rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"]
label_array, color = label_arrays[cls], rainbows[cls] label_array, color = label_arrays[cls], rainbows[cls]
box = [(int(q[0]), int(q[1])), (int(q[2]), int(q[3])), 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]))] (int(q[4]), int(q[5])), (int(q[6]), int(q[7]))]
if bp_np is not None:
if len(bp_np.shape)==1:
bp_np = bp_np[np.newaxis, ...]
for bp in bp_np:
if np.array_equal(bp[:2], np.array([int(q[0]), int(q[1])])):
box.append(bp[-1])
is_new = False is_new = False
if q[11] == 1: if q[11] == 1:
is_new = True is_new = True
if ModelType.CITY_CARPLATE_MODEL.value[1] == str(code) or \
ModelType.CITY_DENSECROWDCOUNT_MODEL.value[1] == str(code) or\
ModelType.CITY_UNDERBUILDCOUNT_MODEL.value[1] == str(code):
box = qs
if cd is None: if cd is None:
det_xywh2[code][cls] = [[cls, box, score, label_array, color, is_new]] det_xywh2[code][cls] = [[cls, box, score, label_array, color, is_new]]
else: else:
det_xywh2[code][cls].append([cls, box, score, label_array, color, is_new]) det_xywh2[code][cls].append(
[cls, box, score, label_array, color, is_new])
if len(det_xywh2) > 0: if len(det_xywh2) > 0:
put_queue(image_queue, (1, [det_xywh2, frame, frame_index_list[i], all_frames, draw_config["font_config"]])) put_queue(image_queue, (1, [det_xywh2, frame, frame_index_list[i], all_frames,
draw_config["font_config"]]))
push_p = push_stream_result.result(timeout=60) push_p = push_stream_result.result(timeout=60)
ai_video_file = write_ai_video_result.result(timeout=60) ai_video_file = write_ai_video_result.result(timeout=60)
@ -264,7 +307,7 @@ class OnPushStreamProcess(PushStreamProcess):
if push_r[0] == 2: if push_r[0] == 2:
logger.info("拉流进程收到控制命令为:{}, requestId: {}",push_r[1] ,request_id) logger.info("拉流进程收到控制命令为:{}, requestId: {}",push_r[1] ,request_id)
if 'algStart' == push_r[1]: self._algStatus = True;logger.info("算法识别开启, requestId: {}", request_id) if 'algStart' == push_r[1]: self._algStatus = True;logger.info("算法识别开启, requestId: {}", request_id)
if 'algStop' == push_r[1]: self._algStatus = False;logger.info("算法识别关闭, requestId: {}", request_id) if 'algStop' == push_r[1]: self._algStatus = False;logger.info("算法识别关闭, requestId: {}", request_id)
if 'stop' == push_r[1]: if 'stop' == push_r[1]:
logger.info("停止推流进程, requestId: {}", request_id) logger.info("停止推流进程, requestId: {}", request_id)
break break
@ -272,7 +315,7 @@ class OnPushStreamProcess(PushStreamProcess):
ex_status = False ex_status = False
logger.info("停止推流进程, requestId: {}", request_id) logger.info("停止推流进程, requestId: {}", request_id)
break break
del push_r del push_r
else: else:
sleep(1) sleep(1)
@ -351,40 +394,72 @@ class OffPushStreamProcess(PushStreamProcess):
# [(2, 操作指令)] 指令操作 # [(2, 操作指令)] 指令操作
if push_r[0] == 1: if push_r[0] == 1:
frame_list, frame_index_list, all_frames, draw_config, push_objs = push_r[1] frame_list, frame_index_list, all_frames, draw_config, push_objs = push_r[1]
# 处理每一帧图片 # 处理每一帧图片
for i, frame in enumerate(frame_list): for i, [frame,mqtt_list] in enumerate(frame_list):
# mqtt传参
border = None
mqttPares = None
if len(mqtt_list) >= 1:
mqttPares = mqtt_list[0]
if mqttPares[0] == 1:
border = mqttPares[1]
pix_dis = int((frame.shape[0]//10)*1.2) pix_dis = int((frame.shape[0]//10)*1.2)
if frame_index_list[i] % 300 == 0 and frame_index_list[i] <= all_frames: if frame_index_list[i] % 300 == 0 and frame_index_list[i] <= all_frames:
task_process = "%.2f" % (float(frame_index_list[i]) / float(all_frames)) task_process = "%.2f" % (float(frame_index_list[i]) / float(all_frames))
put_queue(hb_queue, {"hb_value": task_process}, timeout=2) put_queue(hb_queue, {"hb_value": task_process}, timeout=2)
# 复制帧用来画图 # 复制帧用来画图
copy_frame = frame.copy() copy_frame = frame.copy()
if border is not None:
copy_frame = draw_transparent_red_polygon(copy_frame, np.array(border, np.int32),alpha=0.25)
# 所有问题记录字典 # 所有问题记录字典
det_xywh, thread_p = {}, [] det_xywh, thread_p = {}, []
det_xywh2 = {} det_xywh2 = {'mqttPares':mqttPares}
# 所有问题的矩阵集合 # 所有问题的矩阵集合
qs_np = None qs_np = None
qs_reurn = [] qs_reurn = []
bp_np = None
for det in push_objs[i]: for det in push_objs[i]:
code, det_result = det code, det_result = det
# 每个单独模型处理 # 每个单独模型处理
# 模型编号、100帧的所有问题, 检测目标、颜色、文字图片 # 模型编号、100帧的所有问题, 检测目标、颜色、文字图片
if len(det_result) > 0: if len(det_result) > 0:
font_config, allowedList = draw_config["font_config"], draw_config[code]["allowedList"] font_config, allowedList = draw_config["font_config"], draw_config[code]["allowedList"]
rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"] rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"]
for qs in det_result: for qs in det_result:
box, score, cls = xywh2xyxy2(qs) # 自研车牌模型处理
if ModelType.CITY_CARPLATE_MODEL.value[1] == str(code):
if cls not in allowedList or score < frame_score: cls = 0
continue box = xy2xyxy(qs[1])
label_array, color = label_arrays[cls], rainbows[cls] score = None
if ModelType.CHANNEL2_MODEL.value[1] == str(code) and cls == 2: color = rainbows[cls]
rr = t.submit(draw_name_joint, box, copy_frame, draw_config[code]["label_dict"], score, color, font_config, qs[6]) label_array = None
label_arrays = [None]
rr = t.submit(draw_name_ocr, qs, copy_frame, color)
elif ModelType.CITY_DENSECROWDCOUNT_MODEL.value[1] == str(code) or\
ModelType.CITY_UNDERBUILDCOUNT_MODEL.value[1] == str(code):
cls = 0
box = [(0,0),(0,0),(0,0),(0,0)]
score = None
color = rainbows[cls]
label_array = None
rr = t.submit(draw_name_crowd, qs, copy_frame, color)
else: else:
rr = t.submit(draw_painting_joint, box, copy_frame, label_array, score, color, font_config) box, score, cls = xywh2xyxy2(qs)
if cls not in allowedList or score < frame_score:
continue
if ModelType.CITY_FIREAREA_MODEL.value[1] == str(code):
box.append(qs[-1])
label_array, color = label_arrays[cls], rainbows[cls]
if ModelType.CHANNEL2_MODEL.value[1] == str(code) and cls == 2:
rr = t.submit(draw_name_joint, box, copy_frame, draw_config[code]["label_dict"], score, color, font_config, qs[6])
else:
rr = t.submit(draw_painting_joint, box, copy_frame, label_array, score, color, font_config, border=border)
thread_p.append(rr) thread_p.append(rr)
if det_xywh.get(code) is None: if det_xywh.get(code) is None:
det_xywh[code] = {} det_xywh[code] = {}
cd = det_xywh[code].get(cls) cd = det_xywh[code].get(cls)
@ -392,17 +467,24 @@ class OffPushStreamProcess(PushStreamProcess):
if cd is None: if cd is None:
det_xywh[code][cls] = [[cls, box, score, label_array, color]] det_xywh[code][cls] = [[cls, box, score, label_array, color]]
else: else:
det_xywh[code][cls].append([cls, box, score, label_array, color]) det_xywh[code][cls].append([cls, box, score, label_array, color])
if qs_np is None: if qs_np is None:
qs_np = np.array([box[0][0], box[0][1], box[1][0], box[1][1], qs_np = np.array([box[0][0], box[0][1], box[1][0], box[1][1],
box[2][0], box[2][1], box[3][0], box[3][1], box[2][0], box[2][1], box[3][0], box[3][1],
score, cls, code],dtype=np.float32) score, cls, code],dtype=np.float32)
else: else:
result_li = np.array([box[0][0], box[0][1], box[1][0], box[1][1], result_li = np.array([box[0][0], box[0][1], box[1][0], box[1][1],
box[2][0], box[2][1], box[3][0], box[3][1], box[2][0], box[2][1], box[3][0], box[3][1],
score, cls, code],dtype=np.float32) score, cls, code],dtype=np.float32)
qs_np = np.row_stack((qs_np, result_li)) qs_np = np.row_stack((qs_np, result_li))
if ModelType.CITY_FIREAREA_MODEL.value[1]== str(code):
if bp_np is None:
bp_np = np.array([box[0][0], box[0][1],box[-1]],dtype=object)
else:
bp_li = np.array([box[0][0], box[0][1],box[-1]],dtype=object)
bp_np = np.row_stack((bp_np, bp_li))
if logo: if logo:
frame = add_water_pic(frame, logo, request_id) frame = add_water_pic(frame, logo, request_id)
copy_frame = add_water_pic(copy_frame, logo, request_id) copy_frame = add_water_pic(copy_frame, logo, request_id)
@ -410,11 +492,9 @@ class OffPushStreamProcess(PushStreamProcess):
for r in thread_p: for r in thread_p:
r.result() r.result()
if self._algSwitch and (not self._algStatus): if self._algSwitch and (not self._algStatus):
# frame_merge = video_conjuncing(frame, frame.copy()) frame_merge = video_conjuncing(frame, frame.copy())
frame_merge = frame.copy() else:
else: frame_merge = video_conjuncing(frame, copy_frame)
# frame_merge = video_conjuncing(frame, copy_frame)
frame_merge = copy_frame
# 写识别视频到本地 # 写识别视频到本地
write_ai_video_result = t.submit(write_ai_video, frame_merge, aiFilePath, write_ai_video_result = t.submit(write_ai_video, frame_merge, aiFilePath,
ai_video_file, ai_video_file,
@ -422,10 +502,9 @@ class OffPushStreamProcess(PushStreamProcess):
push_stream_result = t.submit(push_video_stream, frame_merge, push_p, push_url, push_stream_result = t.submit(push_video_stream, frame_merge, push_p, push_url,
p_push_status, request_id) p_push_status, request_id)
if qs_np is not None: if qs_np is not None:
if len(qs_np.shape) == 1: if len(qs_np.shape) == 1:
qs_np = qs_np[np.newaxis,...] qs_np = qs_np[np.newaxis,...]
qs_np_id = qs_np.copy() qs_np_id = qs_np.copy()
b = np.ones(qs_np_id.shape[0]) b = np.ones(qs_np_id.shape[0])
qs_np_id = np.column_stack((qs_np_id,b)) qs_np_id = np.column_stack((qs_np_id,b))
@ -433,7 +512,7 @@ class OffPushStreamProcess(PushStreamProcess):
if picture_similarity: if picture_similarity:
qs_np_tmp = qs_np_id.copy() qs_np_tmp = qs_np_id.copy()
b = np.zeros(qs_np.shape[0]) b = np.zeros(qs_np.shape[0])
qs_reurn = np.column_stack((qs_np,b)) qs_reurn = np.column_stack((qs_np,b))
else: else:
qs_reurn = filterBox(qs_np, qs_np_tmp, pix_dis) qs_reurn = filterBox(qs_np, qs_np_tmp, pix_dis)
if picture_similarity: if picture_similarity:
@ -452,7 +531,7 @@ class OffPushStreamProcess(PushStreamProcess):
if q[11] >= 1: if q[11] >= 1:
cls = int(q[9]) cls = int(q[9])
if not (cls in new_lab): if not (cls in new_lab):
continue # 为了防止其他类别被带出 continue # 为了防止其他类别被带出
code = str(int(q[10])).zfill(3) code = str(int(q[10])).zfill(3)
if det_xywh2.get(code) is None: if det_xywh2.get(code) is None:
det_xywh2[code] = {} det_xywh2[code] = {}
@ -460,16 +539,28 @@ class OffPushStreamProcess(PushStreamProcess):
score = q[8] score = q[8]
rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"] rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"]
label_array, color = label_arrays[cls], rainbows[cls] label_array, color = label_arrays[cls], rainbows[cls]
box = [(int(q[0]), int(q[1])), (int(q[2]), int(q[3])), 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]))] (int(q[4]), int(q[5])), (int(q[6]), int(q[7]))]
if bp_np is not None:
if len(bp_np.shape)==1:
bp_np = bp_np[np.newaxis, ...]
for bp in bp_np:
if np.array_equal(bp[:2], np.array([int(q[0]), int(q[1])])):
box.append(bp[-1])
is_new = False is_new = False
if q[11] == 1: if q[11] == 1:
is_new = True is_new = True
if ModelType.CITY_CARPLATE_MODEL.value[1] == str(code) or \
ModelType.CITY_DENSECROWDCOUNT_MODEL.value[1] == str(code) or\
ModelType.CITY_UNDERBUILDCOUNT_MODEL.value[1] == str(code):
box = qs
if cd is None: if cd is None:
det_xywh2[code][cls] = [[cls, box, score, label_array, color, is_new]] det_xywh2[code][cls] = [[cls, box, score, label_array, color, is_new]]
else: else:
det_xywh2[code][cls].append([cls, box, score, label_array, color, is_new]) det_xywh2[code][cls].append(
if len(det_xywh2) > 0: [cls, box, score, label_array, color, is_new])
if len(det_xywh2) > 1:
put_queue(image_queue, (1, [det_xywh2, frame, frame_index_list[i], all_frames, draw_config["font_config"]])) put_queue(image_queue, (1, [det_xywh2, frame, frame_index_list[i], all_frames, draw_config["font_config"]]))
push_p = push_stream_result.result(timeout=60) push_p = push_stream_result.result(timeout=60)
ai_video_file = write_ai_video_result.result(timeout=60) ai_video_file = write_ai_video_result.result(timeout=60)
@ -477,7 +568,7 @@ class OffPushStreamProcess(PushStreamProcess):
if push_r[0] == 2: if push_r[0] == 2:
logger.info("拉流进程收到控制命令为:{}, requestId: {}",push_r[1] ,request_id) logger.info("拉流进程收到控制命令为:{}, requestId: {}",push_r[1] ,request_id)
if 'algStart' == push_r[1]: self._algStatus = True;logger.info("算法识别开启, requestId: {}", request_id) if 'algStart' == push_r[1]: self._algStatus = True;logger.info("算法识别开启, requestId: {}", request_id)
if 'algStop' == push_r[1]: self._algStatus = False;logger.info("算法识别关闭, requestId: {}", request_id) if 'algStop' == push_r[1]: self._algStatus = False;logger.info("算法识别关闭, requestId: {}", request_id)
if 'stop' == push_r[1]: if 'stop' == push_r[1]:
logger.info("停止推流进程, requestId: {}", request_id) logger.info("停止推流进程, requestId: {}", request_id)
break break

27
config/mqtt/dsp_test_mqtt.yml
View File

@ -1,10 +1,21 @@
mqtt_flag: true mqtt_flag: true
broker : "58.213.148.44" # 业务0为经纬度定位，业务1为入侵算法开关
port : 1883 business: 1
username: "admin" # 经纬度定位
password: "admin##123" location:
#topic: "/topic/v1/airportFly/%s/aiDroneData" broker : "58.213.148.44"
topic: "/topic/v1/airportDrone/THJSQ03B2309TPCTD5QV/realTime/data" port : 1883
# 存储多少条消息到list里 username: "admin"
length: 10 password: "admin##123"
#topic: "/topic/v1/airportFly/%s/aiDroneData"
topic: "/topic/v1/airportDrone/THJSQ03B2309TPCTD5QV/realTime/data"
# 入侵
invade:
broker : "192.168.11.8"
port : 2883
#topic: "/topic/v1/airportFly/%s/aiDroneData"
topic: "test000/topic"
# 存储多少条消息到list里
length: 30

8
config/service/dsp_test_service.yml
View File

@ -25,10 +25,6 @@ service:
# 2 模型追踪 # 2 模型追踪
model_type: 1 model_type: 1
limit: 3 limit: 3
# 新增：大模型识别配置
llm:
server_ip: "192.168.10.11"
api_url: "http://192.168.10.11:8000/generate"
task: task:
# 任务限制5个 # 任务限制5个
limit: 5 limit: 5
@ -37,6 +33,8 @@ service:
#storage source,0--aliyun,1--minio #storage source,0--aliyun,1--minio
storage_source: 0 storage_source: 0
#是否启用mqtt，0--不用，1--启用 #是否启用mqtt，0--不用，1--启用
mqtt_flag: 0 mqtt:
flag: 0
business: 1
#是否启用alg控制功能 #是否启用alg控制功能
algSwitch: False algSwitch: False

768
enums/ModelTypeEnum-jcq.py Normal file
View File

@ -0,0 +1,768 @@
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,'rubCls': 1, 'Rubfilter': 150}
},
'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.5,'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.5,
"iou_thres": 0.5,
"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.25,
"iou_thres": 0.5,
"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, #未挂国旗船只
'uncoverId':5, #未封仓标签
'recScale':1.2,
'target_cls':3.0, #目标种类
'filter_cls':4.0 #被过滤的种类
}},
'models':[
{
#'weight':'../AIlib2/weights/conf/channel2/yolov5.pt',
# 'weight':'../AIlib2/weights/channel2/yolov5_%s_fp16.engine'%(gpuName),
'weight':'/home/thsw2/jcq/test/AIlib2/weights/channel2/best.pt', # yolov5 原来模型基础上增加了未封仓
# '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,
},
} ,
# {
# 'weight':'/home/thsw2/jcq/test/AIlib2/weights1/conf/channel2/yolov5_04.pt', # yolov5_04 添加了uncover 0 4 ；标签 yolov5_jcq
# 'name':'yolov5',
# 'model':yolov5Model,
# 'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.15,'iou_thres':0.25,'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} }
# }
],
'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.5,'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

807
enums/ModelTypeEnum-raw.py Normal file
View File

@ -0,0 +1,807 @@
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),
'modelSize': (1920, 1080),
'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': 10,
'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':5,'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,4,5],'segRegionCnt':1, 'trtFlag_det':False,'trtFlag_seg':False, "score_byClass":{"0":0.8,"1":0.4,"2":0.5,"3":0.5,"4":0.4,"5":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 = ("26", "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
})
SMARTSITE_MODEL = ("28", "028", "智慧工地模型", 'smartSite', lambda device, gpuName: {
'labelnames': [ "工人","塔式起重机","悬臂","起重机","压路机","推土机","挖掘机","卡车","装载机","泵车","混凝土搅拌车","打桩","其他车辆" ],
'postProcess':{'function':default_mix,'pars':{}},
'models':
[
{
'weight':"../AIlib2/weights/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 } },
}
],
'postFile': {
"rainbows": COLOR
},
})
RUBBISH_MODEL = ("29", "029", "垃圾模型", 'rubbish', lambda device, gpuName: {
'labelnames': [ "建筑垃圾","白色垃圾","其他垃圾"],
'postProcess':{'function':default_mix,'pars':{}},
'models':
[
{
'weight':"../AIlib2/weights/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 } },
}
],
'postFile': {
"rainbows": COLOR
},
})
FIREWORK_MODEL = ("30", "030", "烟花模型", 'firework', lambda device, gpuName: {
'labelnames': [ "烟花"],
'postProcess':{'function':default_mix,'pars':{}},
'models':
[
{
'weight':"../AIlib2/weights/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 } },
}
],
'postFile': {
"rainbows": COLOR
},
})
@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

174
enums/ModelTypeEnum.py
View File

@ -14,6 +14,7 @@ from utilsK.drownUtils import mixDrowing_water_postprocess
from utilsK.noParkingUtils import mixNoParking_road_postprocess from utilsK.noParkingUtils import mixNoParking_road_postprocess
from utilsK.illParkingUtils import illParking_postprocess from utilsK.illParkingUtils import illParking_postprocess
from utilsK.pannelpostUtils import pannel_post_process from utilsK.pannelpostUtils import pannel_post_process
from utilsK.securitypostUtils import security_post_process
from stdc import stdcModel from stdc import stdcModel
from yolov5 import yolov5Model from yolov5 import yolov5Model
from p2pNet import p2NnetModel from p2pNet import p2NnetModel
@ -63,7 +64,7 @@ class ModelType(Enum):
"classes": 5, "classes": 5,
"rainbows": COLOR "rainbows": COLOR
}, },
'fiterList':[2],
'Detweights': "../weights/trt/AIlib2/river/yolov5_%s_fp16.engine" % gpuName, 'Detweights': "../weights/trt/AIlib2/river/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/river/stdc_360X640_%s_fp16.engine' % gpuName 'Segweights': '../weights/trt/AIlib2/river/stdc_360X640_%s_fp16.engine' % gpuName
}) })
@ -99,10 +100,8 @@ class ModelType(Enum):
'weight':"../weights/trt/AIlib2/forest2/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径 'weight':"../weights/trt/AIlib2/forest2/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径
'name':'yolov5', 'name':'yolov5',
'model':yolov5Model, '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': { 'postFile': {
@ -112,11 +111,10 @@ class ModelType(Enum):
"classes": 5, "classes": 5,
"rainbows": COLOR "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,###分割模型结果需要保留的等值线数目 'segRegionCnt':2,###分割模型结果需要保留的等值线数目
"pixScale": 1.2, "pixScale": 1.2,
}) })
@ -162,7 +160,8 @@ class ModelType(Enum):
"classes": 10, "classes": 10,
"rainbows": COLOR "rainbows": COLOR
}, },
'allowedList':[0,1,2,3,4,5,6,7,8,9,10,11,12,16,17,18,19,20,21,22], '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, 'Detweights': "../weights/trt/AIlib2/highWay2/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/highWay2/stdc_360X640_%s_fp16.engine' % gpuName 'Segweights': '../weights/trt/AIlib2/highWay2/stdc_360X640_%s_fp16.engine' % gpuName
}) })
@ -231,7 +230,7 @@ class ModelType(Enum):
"classes": 5, "classes": 5,
"rainbows": COLOR "rainbows": COLOR
}, },
'Segweights': None 'Segweights': None,
}) })
ANGLERSWIMMER_MODEL = ("9", "009", "钓鱼游泳模型", 'AnglerSwimmer', lambda device, gpuName: { ANGLERSWIMMER_MODEL = ("9", "009", "钓鱼游泳模型", 'AnglerSwimmer', lambda device, gpuName: {
@ -345,7 +344,8 @@ class ModelType(Enum):
'function': riverDetSegMixProcess, 'function': riverDetSegMixProcess,
'pars': { 'pars': {
'slopeIndex': [1, 3, 4, 7], 'slopeIndex': [1, 3, 4, 7],
'riverIou': 0.1 'riverIou': 0.1,
'scale': 0.25
} }
} }
}, },
@ -374,13 +374,14 @@ class ModelType(Enum):
}, },
'models':[ 'models':[
{ {
'weight':'../weights/trt/AIlib2/cityMangement3/yolov5_%s_fp16.engine'%(gpuName), 'weight':'../weights/trt/AIlib2/cityMangement3/yolov5_%s_fp16.engine'%(gpuName),
'name':'yolov5', 'name':'yolov5',
'model':yolov5Model, '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/pth/AIlib2/cityMangement3/dmpr.pth', 'weight':'../weights/trt/AIlib2/cityMangement3/dmpr_3090.engine',
#'weight':'../weights/pth/AIlib2/cityMangement3/dmpr.pth',
'par':{ 'par':{
'depth_factor':32,'NUM_FEATURE_MAP_CHANNEL':6,'dmpr_thresh':0.1, 'dmprimg_size':640, 'depth_factor':32,'NUM_FEATURE_MAP_CHANNEL':6,'dmpr_thresh':0.1, 'dmprimg_size':640,
'name':'dmpr' 'name':'dmpr'
@ -403,7 +404,7 @@ class ModelType(Enum):
"classes": 8, "classes": 8,
"rainbows": COLOR "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,###分割模型结果需要保留的等值线数目 'segRegionCnt':2,###分割模型结果需要保留的等值线数目
"pixScale": 1.2, "pixScale": 1.2,
}) })
@ -568,10 +569,10 @@ class ModelType(Enum):
'weight':'../weights/trt/AIlib2/channel2/yolov5_%s_fp16.engine'%(gpuName), 'weight':'../weights/trt/AIlib2/channel2/yolov5_%s_fp16.engine'%(gpuName),
'name':'yolov5', 'name':'yolov5',
'model':yolov5Model, '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', 'name':'ocr',
'model':ocrModel, 'model':ocrModel,
'par':{ 'par':{
@ -587,7 +588,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, 'segPar': None,
'postFile': { 'postFile': {
"name": "post_process", "name": "post_process",
@ -597,6 +597,8 @@ class ModelType(Enum):
"rainbows": COLOR "rainbows": COLOR
}, },
'Segweights': None, '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: { RIVERT_MODEL = ("25", "025", "河道检测模型(T)", 'riverT', lambda device, gpuName: {
@ -633,8 +635,6 @@ class ModelType(Enum):
'Segweights': '../weights/trt/AIlib2/riverT/stdc_360X640_%s_fp16.engine' % gpuName 'Segweights': '../weights/trt/AIlib2/riverT/stdc_360X640_%s_fp16.engine' % gpuName
}) })
FORESTCROWD_FARM_MODEL = ("26", "026", "森林人群模型", 'forestCrowd', lambda device, gpuName: { FORESTCROWD_FARM_MODEL = ("26", "026", "森林人群模型", 'forestCrowd', lambda device, gpuName: {
'labelnames': ["林斑", "病死树", "行人", "火焰", "烟雾","人群"], 'labelnames': ["林斑", "病死树", "行人", "火焰", "烟雾","人群"],
'postProcess':{'function':gather_post_process,'pars':{'pedestrianId':2,'crowdThreshold':4,'gatherId':5,'distancePersonScale':2.0}}, 'postProcess':{'function':gather_post_process,'pars':{'pedestrianId':2,'crowdThreshold':4,'gatherId':5,'distancePersonScale':2.0}},
@ -644,7 +644,7 @@ class ModelType(Enum):
'weight':"../weights/trt/AIlib2/forestCrowd/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径 'weight':"../weights/trt/AIlib2/forestCrowd/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径
'name':'yolov5', 'name':'yolov5',
'model':yolov5Model, '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},
} }
@ -657,7 +657,7 @@ class ModelType(Enum):
"classes": 5, "classes": 5,
"rainbows": COLOR "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,###分割模型结果需要保留的等值线数目 'segRegionCnt':2,###分割模型结果需要保留的等值线数目
"pixScale": 1.2, "pixScale": 1.2,
@ -705,6 +705,7 @@ class ModelType(Enum):
"classes": 10, "classes": 10,
"rainbows": COLOR "rainbows": COLOR
}, },
'fiterltList': [11,12,13,14,15,16,17],
'Detweights': "../weights/trt/AIlib2/highWay2T/yolov5_%s_fp16.engine" % gpuName, 'Detweights': "../weights/trt/AIlib2/highWay2T/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/highWay2T/stdc_360X640_%s_fp16.engine' % gpuName 'Segweights': '../weights/trt/AIlib2/highWay2T/stdc_360X640_%s_fp16.engine' % gpuName
}) })
@ -718,7 +719,7 @@ class ModelType(Enum):
'weight':"../weights/trt/AIlib2/smartSite/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径 'weight':"../weights/trt/AIlib2/smartSite/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径
'name':'yolov5', 'name':'yolov5',
'model':yolov5Model, '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},
} }
@ -726,6 +727,7 @@ class ModelType(Enum):
'postFile': { 'postFile': {
"rainbows": COLOR "rainbows": COLOR
}, },
"score_byClass": {0: 0.25, 1: 0.3, 2: 0.3, 3: 0.3}
}) })
@ -738,7 +740,7 @@ class ModelType(Enum):
'weight':"../weights/trt/AIlib2/rubbish/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径 'weight':"../weights/trt/AIlib2/rubbish/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径
'name':'yolov5', 'name':'yolov5',
'model':yolov5Model, '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},
} }
@ -746,6 +748,7 @@ class ModelType(Enum):
'postFile': { 'postFile': {
"rainbows": COLOR "rainbows": COLOR
}, },
"score_byClass": {0: 0.25, 1: 0.3, 2: 0.3, 3: 0.3}
}) })
@ -758,7 +761,7 @@ class ModelType(Enum):
'weight':"../weights/trt/AIlib2/firework/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径 'weight':"../weights/trt/AIlib2/firework/yolov5_%s_fp16.engine"%(gpuName),###检测模型路径
'name':'yolov5', 'name':'yolov5',
'model':yolov5Model, '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 },
} }
@ -766,7 +769,6 @@ class ModelType(Enum):
'postFile': { 'postFile': {
"rainbows": COLOR "rainbows": COLOR
}, },
}) })
TRAFFIC_SPILL_MODEL = ("50", "501", "高速公路抛洒物模型", 'highWaySpill', lambda device, gpuName: { TRAFFIC_SPILL_MODEL = ("50", "501", "高速公路抛洒物模型", 'highWaySpill', lambda device, gpuName: {
@ -808,7 +810,7 @@ class ModelType(Enum):
"classes": 2, "classes": 2,
"rainbows": COLOR "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, 'Detweights': "../weights/trt/AIlib2/highWaySpill/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/highWaySpill/stdc_360X640_%s_fp16.engine' % gpuName 'Segweights': '../weights/trt/AIlib2/highWaySpill/stdc_360X640_%s_fp16.engine' % gpuName
@ -853,7 +855,7 @@ class ModelType(Enum):
"classes": 4, "classes": 4,
"rainbows": COLOR "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, 'Detweights': "../weights/trt/AIlib2/highWayCthc/yolov5_%s_fp16.engine" % gpuName,
'Segweights': '../weights/trt/AIlib2/highWayCthc/stdc_360X640_%s_fp16.engine' % gpuName 'Segweights': '../weights/trt/AIlib2/highWayCthc/stdc_360X640_%s_fp16.engine' % gpuName
@ -869,14 +871,15 @@ class ModelType(Enum):
'name': 'yolov5', 'name': 'yolov5',
'model': yolov5Model, 'model': yolov5Model,
'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.25, 'iou_thres': 0.45, 'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.25, 'iou_thres': 0.45,
'allowedList': [0,1,2], 'segRegionCnt': 1, 'trtFlag_det': True, 'segRegionCnt': 1, 'trtFlag_det': True,
'trtFlag_seg': False, "score_byClass": {"0": 0.25, "1": 0.3, "2": 0.3, "3": 0.3}}, 'trtFlag_seg': False},
} }
], ],
'postFile': { 'postFile': {
"rainbows": COLOR "rainbows": COLOR
}, },
'fiterList':[0]
}) })
@ -886,25 +889,27 @@ class ModelType(Enum):
'rainbows': COLOR, 'rainbows': COLOR,
'models': [ 'models': [
{ {
'trtFlag_det': False, #'weight': '../weights/pth/AIlib2/carplate/plate_yolov5s_v3.jit',
'weight': '../weights/pth/AIlib2/carplate/plate_yolov5s_v3.jit', 'weight': '../weights/trt/AIlib2/carplate/yolov5_%s_fp16.engine' % (gpuName),
'name': 'yolov5', 'name': 'yolov5',
'model': yolov5Model, 'model': yolov5Model,
'par': { 'par': {
'trtFlag_det': True,
'device': 'cuda:0', 'device': 'cuda:0',
'half': False, 'half': True,
'conf_thres': 0.4, 'conf_thres': 0.4,
'iou_thres': 0.45, 'iou_thres': 0.45,
'nc': 1, 'nc': 1,
'plate':8,
'plate_dilate': (0.5, 0.1) 'plate_dilate': (0.5, 0.1)
}, },
}, },
{ {
'trtFlag_ocr': False, 'weight' : '../weights/trt/AIlib2/ocr2/crnn_ch_%s_fp16_192X32.engine'%(gpuName),
'weight': '../weights/pth/AIlib2/ocr2/crnn_ch.pth',
'name': 'ocr', 'name': 'ocr',
'model': ocrModel, 'model': ocrModel,
'par': { 'par': {
'trtFlag_ocr': True,
'char_file': '../AIlib2/conf/ocr2/benchmark.txt', 'char_file': '../AIlib2/conf/ocr2/benchmark.txt',
'mode': 'ch', 'mode': 'ch',
'nc': 3, 'nc': 3,
@ -928,10 +933,8 @@ class ModelType(Enum):
'name': 'yolov5', 'name': 'yolov5',
'model': yolov5Model, 'model': yolov5Model,
'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.50, 'iou_thres': 0.45, 'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.50, 'iou_thres': 0.45,
'allowedList': list(range(20)), 'segRegionCnt': 1, 'trtFlag_det': True, 'segRegionCnt': 1, 'trtFlag_det': True,'trtFlag_seg': False},
'trtFlag_seg': False, "score_byClass": {"0": 0.50, "1": 0.3, "2": 0.3, "3": 0.3}},
} }
], ],
'postFile': { 'postFile': {
"rainbows": COLOR "rainbows": COLOR
@ -949,8 +952,7 @@ class ModelType(Enum):
'name': 'yolov5', 'name': 'yolov5',
'model': yolov5Model, 'model': yolov5Model,
'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.50, 'iou_thres': 0.45, 'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.50, 'iou_thres': 0.45,
'allowedList': list(range(20)), 'segRegionCnt': 1, 'trtFlag_det': True, 'segRegionCnt': 1, 'trtFlag_det': True, 'trtFlag_seg': False},
'trtFlag_seg': False, "score_byClass": {"0": 0.50, "1": 0.3, "2": 0.3, "3": 0.3}},
} }
], ],
@ -976,12 +978,14 @@ class ModelType(Enum):
'row': 2, 'row': 2,
'line': 2, 'line': 2,
'point_loss_coef': 0.45, 'point_loss_coef': 0.45,
'conf': 0.25, 'conf': 0.65,
'gpu_id': 0, 'gpu_id': 0,
'eos_coef': '0.5', 'eos_coef': '0.5',
'set_cost_class': 1, 'set_cost_class': 1,
'set_cost_point': 0.05, 'set_cost_point': 0.05,
'backbone': 'vgg16_bn' 'backbone': 'vgg16_bn',
'expend': 10,
'psize': 2,
}, },
}], }],
}) })
@ -995,8 +999,7 @@ class ModelType(Enum):
'name': 'yolov5', 'name': 'yolov5',
'model': yolov5Model, 'model': yolov5Model,
'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.50, 'iou_thres': 0.45, 'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.50, 'iou_thres': 0.45,
'allowedList': list(range(20)), 'segRegionCnt': 1, 'trtFlag_det': True, 'segRegionCnt': 1, 'trtFlag_det': True, 'trtFlag_seg': False},
'trtFlag_seg': False, "score_byClass": {"0": 0.50, "1": 0.3, "2": 0.3, "3": 0.3}},
} }
], ],
@ -1006,6 +1009,89 @@ class ModelType(Enum):
}) })
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, # 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}
})
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 @staticmethod
def checkCode(code): def checkCode(code):
for model in ModelType: for model in ModelType:

95
util/ModelUtils.py
View File

@ -27,6 +27,7 @@ import torch
import tensorrt as trt import tensorrt as trt
from utilsK.jkmUtils import pre_process, post_process, get_return_data from utilsK.jkmUtils import pre_process, post_process, get_return_data
from DMPR import DMPRModel from DMPR import DMPRModel
from segment_anything import SamPredictor, sam_model_registry
FONT_PATH = "../AIlib2/conf/platech.ttf" FONT_PATH = "../AIlib2/conf/platech.ttf"
@ -36,6 +37,7 @@ class OneModel:
def __init__(self, device, allowedList=None, requestId=None, modeType=None, gpu_name=None, base_dir=None, env=None): def __init__(self, device, allowedList=None, requestId=None, modeType=None, gpu_name=None, base_dir=None, env=None):
try: try:
start = time.time()
logger.info("########################加载{}########################, requestId:{}", modeType.value[2], logger.info("########################加载{}########################, requestId:{}", modeType.value[2],
requestId) requestId)
par = modeType.value[4](str(device), gpu_name) par = modeType.value[4](str(device), gpu_name)
@ -46,8 +48,12 @@ class OneModel:
new_device = select_device(par.get('device')) new_device = select_device(par.get('device'))
half = new_device.type != 'cpu' half = new_device.type != 'cpu'
Detweights = par['Detweights'] Detweights = par['Detweights']
with open(Detweights, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.ERROR)) as runtime: if par['trtFlag_det']:
model = runtime.deserialize_cuda_engine(f.read()) with open(Detweights, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.ERROR)) as runtime:
model = runtime.deserialize_cuda_engine(f.read())
else:
model = attempt_load(Detweights, map_location=new_device) # load FP32 model
if half: model.half()
par['segPar']['seg_nclass'] = par['seg_nclass'] par['segPar']['seg_nclass'] = par['seg_nclass']
Segweights = par['Segweights'] Segweights = par['Segweights']
if Segweights: if Segweights:
@ -64,10 +70,11 @@ class OneModel:
'ovlap_thres_crossCategory': postFile.get("ovlap_thres_crossCategory"), 'ovlap_thres_crossCategory': postFile.get("ovlap_thres_crossCategory"),
'iou_thres': postFile["iou_thres"], 'iou_thres': postFile["iou_thres"],
# 对高速模型进行过滤 # 对高速模型进行过滤
'allowedList': par['allowedList'] if modeType.value[0] == '3' else [],
'segRegionCnt': par['segRegionCnt'], 'segRegionCnt': par['segRegionCnt'],
'trtFlag_det': par['trtFlag_det'], 'trtFlag_det': par['trtFlag_det'],
'trtFlag_seg': par['trtFlag_seg'] 'trtFlag_seg': par['trtFlag_seg'],
'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_param = {
"model": model, "model": model,
@ -82,6 +89,7 @@ class OneModel:
logger.error("模型加载异常：{}, requestId:{}", format_exc(), requestId) logger.error("模型加载异常：{}, requestId:{}", format_exc(), requestId)
raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0], raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
ExceptionType.MODEL_LOADING_EXCEPTION.value[1]) ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
logger.info("模型初始化时间：{}, requestId:{}", time.time() - start, requestId)
# 纯分类模型 # 纯分类模型
class cityManagementModel: class cityManagementModel:
__slots__ = "model_conf" __slots__ = "model_conf"
@ -99,6 +107,8 @@ class cityManagementModel:
model_param = { model_param = {
"modelList": modelList, "modelList": modelList,
"postProcess": postProcess, "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) self.model_conf = (modeType, model_param, allowedList, names, rainbows)
except Exception: except Exception:
@ -107,15 +117,14 @@ class cityManagementModel:
ExceptionType.MODEL_LOADING_EXCEPTION.value[1]) ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
def detSeg_demo2(args): def detSeg_demo2(args):
model_conf, frame, request_id = 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: 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 return result
except ServiceException as s: except ServiceException as s:
raise s raise s
except Exception: 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) logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), request_id)
raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0], raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0],
ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1]) ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1])
@ -123,11 +132,6 @@ def detSeg_demo2(args):
def model_process(args): def model_process(args):
model_conf, frame, request_id = args model_conf, frame, request_id = args
model_param, names, rainbows = model_conf[1], model_conf[3], model_conf[4] 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: try:
return AI_process([frame], model_param['model'], model_param['segmodel'], names, model_param['label_arraylist'], return AI_process([frame], model_param['model'], model_param['segmodel'], names, model_param['label_arraylist'],
rainbows, objectPar=model_param['objectPar'], font=model_param['digitFont'], rainbows, objectPar=model_param['objectPar'], font=model_param['digitFont'],
@ -160,7 +164,13 @@ class TwoModel:
Detweights = par['Detweights'] Detweights = par['Detweights']
with open(Detweights, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.ERROR)) as runtime: with open(Detweights, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.ERROR)) as runtime:
model = runtime.deserialize_cuda_engine(f.read()) 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'] postFile = par['postFile']
conf_thres = postFile["conf_thres"] conf_thres = postFile["conf_thres"]
iou_thres = postFile["iou_thres"] iou_thres = postFile["iou_thres"]
@ -174,7 +184,10 @@ class TwoModel:
"conf_thres": conf_thres, "conf_thres": conf_thres,
"iou_thres": iou_thres, "iou_thres": iou_thres,
"trtFlag_det": par['trtFlag_det'], "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) self.model_conf = (modeType, model_param, allowedList, names, rainbows)
except Exception: except Exception:
@ -182,16 +195,15 @@ class TwoModel:
raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0], raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
ExceptionType.MODEL_LOADING_EXCEPTION.value[1]) ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
logger.info("模型初始化时间：{}, requestId:{}", time.time() - s, requestId) logger.info("模型初始化时间：{}, requestId:{}", time.time() - s, requestId)
def forest_process(args): def forest_process(args):
model_conf, frame, request_id = args model_conf, frame, request_id = args
model_param, names, rainbows = model_conf[1], model_conf[3], model_conf[4] model_param, names, rainbows = model_conf[1], model_conf[3], model_conf[4]
try: try:
return AI_process_forest([frame], model_param['model'], model_param['segmodel'], names, 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['label_arraylist'], rainbows, model_param['half'], model_param['device'],
model_param['conf_thres'], model_param['iou_thres'], [], font=model_param['digitFont'], model_param['conf_thres'], model_param['iou_thres'],font=model_param['digitFont'],
trtFlag_det=model_param['trtFlag_det'], SecNms=model_param['otc']) 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: except ServiceException as s:
raise s raise s
except Exception: except Exception:
@ -200,7 +212,6 @@ def forest_process(args):
logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), request_id) logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), request_id)
raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0], raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0],
ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1]) ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1])
class MultiModel: class MultiModel:
__slots__ = "model_conf" __slots__ = "model_conf"
@ -219,6 +230,8 @@ class MultiModel:
model_param = { model_param = {
"modelList": modelList, "modelList": modelList,
"postProcess": postProcess, "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) self.model_conf = (modeType, model_param, allowedList, names, rainbows)
except Exception: except Exception:
@ -226,13 +239,13 @@ class MultiModel:
raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0], raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
ExceptionType.MODEL_LOADING_EXCEPTION.value[1]) ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
logger.info("模型初始化时间：{}, requestId:{}", time.time() - s, requestId) logger.info("模型初始化时间：{}, requestId:{}", time.time() - s, requestId)
def channel2_process(args): def channel2_process(args):
model_conf, frame, request_id = 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: try:
start = time.time() 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)) # print("AI_process_C use time = {}".format(time.time()-start))
return result return result
except ServiceException as s: except ServiceException as s:
@ -241,7 +254,6 @@ def channel2_process(args):
logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), request_id) logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), request_id)
raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0], raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0],
ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1]) ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1])
def get_label_arraylist(*args): def get_label_arraylist(*args):
width, height, names, rainbows = args width, height, names, rainbows = args
# line = int(round(0.002 * (height + width) / 2) + 1) # line = int(round(0.002 * (height + width) / 2) + 1)
@ -262,8 +274,6 @@ def get_label_arraylist(*args):
'label_location': 'leftTop'} 'label_location': 'leftTop'}
label_arraylist = get_label_arrays(names, rainbows, fontSize=text_height, fontPath=FONT_PATH) label_arraylist = get_label_arrays(names, rainbows, fontSize=text_height, fontPath=FONT_PATH)
return digitFont, label_arraylist, (line, text_width, text_height, fontScale, tf) return digitFont, label_arraylist, (line, text_width, text_height, fontScale, tf)
# 船只模型 # 船只模型
class ShipModel: class ShipModel:
__slots__ = "model_conf" __slots__ = "model_conf"
@ -289,8 +299,6 @@ class ShipModel:
raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0], raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
ExceptionType.MODEL_LOADING_EXCEPTION.value[1]) ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
logger.info("模型初始化时间：{}, requestId:{}", time.time() - s, requestId) logger.info("模型初始化时间：{}, requestId:{}", time.time() - s, requestId)
def obb_process(args): def obb_process(args):
model_conf, frame, request_id = args model_conf, frame, request_id = args
model_param = model_conf[1] model_param = model_conf[1]
@ -305,7 +313,6 @@ def obb_process(args):
logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), request_id) logger.error("算法模型分析异常：{}, requestId:{}", format_exc(), request_id)
raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0], raise ServiceException(ExceptionType.MODEL_ANALYSE_EXCEPTION.value[0],
ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1]) ExceptionType.MODEL_ANALYSE_EXCEPTION.value[1])
# 车牌分割模型、健康码、行程码分割模型 # 车牌分割模型、健康码、行程码分割模型
class IMModel: class IMModel:
__slots__ = "model_conf" __slots__ = "model_conf"
@ -329,7 +336,7 @@ class IMModel:
new_device = torch.device(par['device']) new_device = torch.device(par['device'])
model = torch.jit.load(par[img_type]['weights']) model = torch.jit.load(par[img_type]['weights'])
logger.info("########################加载 jit 模型成功 成功 ########################, requestId:{}", logger.info("########################加载 jit 模型成功 成功 ########################, requestId:{}",
requestId) requestId)
self.model_conf = (modeType, allowedList, new_device, model, par, img_type) self.model_conf = (modeType, allowedList, new_device, model, par, img_type)
except Exception: except Exception:
@ -383,15 +390,12 @@ class CARPLATEModel:
detpar = par['models'][0]['par'] detpar = par['models'][0]['par']
# new_device = torch.device(par['device']) # new_device = torch.device(par['device'])
# modelList=[ modelPar['model'](weights=modelPar['weight'],par=modelPar['par']) for modelPar in par['models'] ] # modelList=[ modelPar['model'](weights=modelPar['weight'],par=modelPar['par']) for modelPar in par['models'] ]
logger.info("########################加载 plate_yolov5s_v3.jit 成功 ########################, requestId:{}",
requestId)
self.model_conf = (modeType, device, modelList, detpar, par['rainbows']) self.model_conf = (modeType, device, modelList, detpar, par['rainbows'])
except Exception: except Exception:
logger.error("模型加载异常：{}, requestId:{}", format_exc(), requestId) logger.error("模型加载异常：{}, requestId:{}", format_exc(), requestId)
raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0], raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
ExceptionType.MODEL_LOADING_EXCEPTION.value[1]) ExceptionType.MODEL_LOADING_EXCEPTION.value[1])
class DENSECROWDCOUNTModel: class DENSECROWDCOUNTModel:
__slots__ = "model_conf" __slots__ = "model_conf"
@ -402,8 +406,8 @@ class DENSECROWDCOUNTModel:
par = modeType.value[4](str(device), gpu_name) par = modeType.value[4](str(device), gpu_name)
rainbows = par["rainbows"] rainbows = par["rainbows"]
models=[ modelPar['model'](weights=modelPar['weight'],par=modelPar['par']) for modelPar in par['models'] ] models=[ modelPar['model'](weights=modelPar['weight'],par=modelPar['par']) for modelPar in par['models'] ]
postPar = par['models'][0]['par'] postPar = [pp['par'] for pp in par['models']]
self.model_conf = (modeType, device, models[0], postPar, rainbows) self.model_conf = (modeType, device, models, postPar, rainbows)
except Exception: except Exception:
logger.error("模型加载异常：{}, requestId:{}", format_exc(), requestId) logger.error("模型加载异常：{}, requestId:{}", format_exc(), requestId)
raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0], raise ServiceException(ExceptionType.MODEL_LOADING_EXCEPTION.value[0],
@ -752,4 +756,25 @@ MODEL_CONFIG = {
None, None,
lambda x: cc_process(x) lambda x: cc_process(x)
), ),
# 加载建筑物下行人检测模型
ModelType.CITY_UNDERBUILDCOUNT_MODEL.value[1]: (
lambda x, y, r, t, z, h: DENSECROWDCOUNTModel(x, y, r, ModelType.CITY_UNDERBUILDCOUNT_MODEL, t, z, h),
ModelType.CITY_UNDERBUILDCOUNT_MODEL,
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)
),
# 加载安防模型
ModelType.CITY_SECURITY_MODEL.value[1]: (
lambda x, y, r, t, z, h: cityManagementModel(x, y, r, ModelType.CITY_SECURITY_MODEL, t, z, h),
ModelType.CITY_SECURITY_MODEL,
lambda x, y, z: one_label(x, y, z),
lambda x: detSeg_demo2(x)
),
} }

254
util/PlotsUtils.py
View File

@ -2,9 +2,10 @@ import cv2
import numpy as np import numpy as np
from PIL import Image, ImageDraw, ImageFont from PIL import Image, ImageDraw, ImageFont
import unicodedata import unicodedata
from loguru import logger
FONT_PATH = "../AIlib2/conf/platech.ttf" FONT_PATH = "../AIlib2/conf/platech.ttf"
zhFont = ImageFont.truetype(FONT_PATH, 20, encoding="utf-8") zhFont = ImageFont.truetype(FONT_PATH, 20, encoding="utf-8")
def get_label_array(color=None, label=None, font=None, fontSize=40, unify=False): def get_label_array(color=None, label=None, font=None, fontSize=40, unify=False):
if unify: if unify:
@ -23,7 +24,6 @@ def get_label_array(color=None, label=None, font=None, fontSize=40, unify=False)
im_array = cv2.resize(im_array, (0, 0), fx=scale, fy=scale) im_array = cv2.resize(im_array, (0, 0), fx=scale, fy=scale)
return im_array return im_array
def get_label_arrays(labelNames, colors, fontSize=40, fontPath="platech.ttf"): def get_label_arrays(labelNames, colors, fontSize=40, fontPath="platech.ttf"):
font = ImageFont.truetype(fontPath, fontSize, encoding='utf-8') font = ImageFont.truetype(fontPath, fontSize, encoding='utf-8')
label_arraylist = [get_label_array(colors[i % 20], label_name, font, fontSize) for i, label_name in label_arraylist = [get_label_array(colors[i % 20], label_name, font, fontSize) for i, label_name in
@ -49,6 +49,48 @@ def get_label_array_dict(colors, fontSize=40, fontPath="platech.ttf"):
zh_dict[code] = arr zh_dict[code] = arr
return zh_dict return zh_dict
def get_label_left(x0,y1,label_array,img):
imh, imw = img.shape[0:2]
lh, lw = label_array.shape[0:2]
# x1 框框左上x位置 + 描述的宽
# y0 框框左上y位置 - 描述的高
x1, y0 = x0 + lw, y1 - lh
# 如果y0小于0, 说明超过上边框
if y0 < 0:
y0 = 0
# y1等于文字高度
y1 = y0 + lh
# 如果y1框框的高大于图片高度
if y1 > imh:
# y1等于图片高度
y1 = imh
# y0等于y1减去文字高度
y0 = y1 - lh
# 如果x0小于0
if x0 < 0:
x0 = 0
x1 = x0 + lw
if x1 > imw:
x1 = imw
x0 = x1 - lw
return x0,y0,x1,y1
def get_label_right(x1,y0,label_array):
lh, lw = label_array.shape[0:2]
# x1 框框右上x位置 + 描述的宽
# y0 框框右上y位置 - 描述的高
x0, y1 = x1 - lw, y0 - lh
# 如果y0小于0, 说明超过上边框
if y0 < 0 or y1 < 0:
y1 = 0
# y1等于文字高度
y0 = y1 + lh
# 如果x0小于0
if x0 < 0 or x1 < 0:
x0 = 0
x1 = x0 + lw
return x0,y1,x1,y0
def xywh2xyxy(box): def xywh2xyxy(box):
if not isinstance(box[0], (list, tuple, np.ndarray)): if not isinstance(box[0], (list, tuple, np.ndarray)):
@ -74,42 +116,24 @@ def xy2xyxy(box):
box = [(x1, y1), (x2, y1), (x2, y2), (x1, y2)] box = [(x1, y1), (x2, y1), (x2, y2), (x1, y2)]
return box return box
def draw_painting_joint(box, img, label_array, score=0.5, color=None, config=None, isNew=False): def draw_painting_joint(box, img, label_array, score=0.5, color=None, config=None, isNew=False, border=None):
# 识别问题描述图片的高、宽 # 识别问题描述图片的高、宽
lh, lw = label_array.shape[0:2]
# 图片的长度和宽度 # 图片的长度和宽度
imh, imw = img.shape[0:2] if border is not None:
border = np.array(border,np.int32)
color,label_array=draw_name_border(box,color,label_array,border)
#img = draw_transparent_red_polygon(img,border,'',alpha=0.1)
lh, lw = label_array.shape[0:2]
tl = config[0]
if isinstance(box[-1], np.ndarray):
return draw_name_points(img,box,color)
label = ' %.2f' % score
box = xywh2xyxy(box) box = xywh2xyxy(box)
# 框框左上的位置 # 框框左上的位置
x0, y1 = box[0][0], box[0][1] x0, y1 = box[0][0], box[0][1]
# if score_location == 'leftTop': x0, y0, x1, y1 = get_label_left(x0, y1, label_array, img)
# x0, y1 = box[0][0], box[0][1]
# # 框框左下的位置
# elif score_location == 'leftBottom':
# x0, y1 = box[3][0], box[3][1]
# else:
# x0, y1 = box[0][0], box[0][1]
# x1 框框左上x位置 + 描述的宽
# y0 框框左上y位置 - 描述的高
x1, y0 = x0 + lw, y1 - lh
# 如果y0小于0, 说明超过上边框
if y0 < 0:
y0 = 0
# y1等于文字高度
y1 = y0 + lh
# 如果y1框框的高大于图片高度
if y1 > imh:
# y1等于图片高度
y1 = imh
# y0等于y1减去文字高度
y0 = y1 - lh
# 如果x0小于0
if x0 < 0:
x0 = 0
x1 = x0 + lw
if x1 > imw:
x1 = imw
x0 = x1 - lw
# box_tl = max(int(round(imw / 1920 * 3)), 1) or round(0.002 * (imh + imw) / 2) + 1 # box_tl = max(int(round(imw / 1920 * 3)), 1) or round(0.002 * (imh + imw) / 2) + 1
''' '''
1. img（array） 为ndarray类型（可以为cv.imread）直接读取的数据 1. img（array） 为ndarray类型（可以为cv.imread）直接读取的数据
@ -119,14 +143,12 @@ def draw_painting_joint(box, img, label_array, score=0.5, color=None, config=Non
5. thickness（int）：画线的粗细 5. thickness（int）：画线的粗细
6. shift：顶点坐标中小数的位数 6. shift：顶点坐标中小数的位数
''' '''
tl = config[0] img[y0:y1, x0:x1, :] = label_array
box1 = np.asarray(box, np.int32) box1 = np.asarray(box, np.int32)
cv2.polylines(img, [box1], True, color, tl) cv2.polylines(img, [box1], True, color, tl)
img[y0:y1, x0:x1, :] = label_array
pts_cls = [(x0, y0), (x1, y1)] pts_cls = [(x0, y0), (x1, y1)]
# 把英文字符score画到类别旁边 # 把英文字符score画到类别旁边
# tl = max(int(round(imw / 1920 * 3)), 1) or round(0.002 * (imh + imw) / 2) + 1 # tl = max(int(round(imw / 1920 * 3)), 1) or round(0.002 * (imh + imw) / 2) + 1
label = ' %.2f' % score
# tf = max(tl, 1) # tf = max(tl, 1)
# fontScale = float(format(imw / 1920 * 1.1, '.2f')) or tl * 0.33 # fontScale = float(format(imw / 1920 * 1.1, '.2f')) or tl * 0.33
# fontScale = tl * 0.33 # fontScale = tl * 0.33
@ -224,12 +246,11 @@ def draw_name_joint(box, img, label_array_dict, score=0.5, color=None, config=No
cv2.putText(img, label, p3, 0, config[3], [225, 255, 255], thickness=config[4], lineType=cv2.LINE_AA) cv2.putText(img, label, p3, 0, config[3], [225, 255, 255], thickness=config[4], lineType=cv2.LINE_AA)
return img, box return img, box
def draw_name_ocr(box, img, color, label, line_thickness=2, outfontsize=40): def draw_name_ocr(box, img, color, line_thickness=2, outfontsize=40):
font = ImageFont.truetype(FONT_PATH, outfontsize, encoding='utf-8') font = ImageFont.truetype(FONT_PATH, outfontsize, encoding='utf-8')
#(color=None, label=None, font=None, fontSize=40, unify=False) # (color=None, label=None, font=None, fontSize=40, unify=False)
label_zh = get_label_array(color, label, font, outfontsize) label_zh = get_label_array(color, box[0], font, outfontsize)
return plot_one_box_auto(box, img, color, line_thickness, label_zh) return plot_one_box_auto(box[1], img, color, line_thickness, label_zh)
def filterBox(det0, det1, pix_dis): def filterBox(det0, det1, pix_dis):
# det0为 (m1, 11) 矩阵 # det0为 (m1, 11) 矩阵
# det1为 (m2, 12) 矩阵 # det1为 (m2, 12) 矩阵
@ -262,7 +283,7 @@ def filterBox(det0, det1, pix_dis):
x_c, y_c = (x3+x4)//2, (y3+y4)//2 x_c, y_c = (x3+x4)//2, (y3+y4)//2
dis = (x2_c - x_c)**2 + (y2_c - y_c)**2 dis = (x2_c - x_c)**2 + (y2_c - y_c)**2
mask = (joint_det[..., 9] == joint_det[..., 20]) & (dis <= pix_dis**2) mask = (joint_det[..., 9] == joint_det[..., 20]) & (dis <= pix_dis**2)
# 类别相同 & 中心点在上一帧的框内 判断为True # 类别相同 & 中心点在上一帧的框内 判断为True
res = np.sum(mask, axis=1) res = np.sum(mask, axis=1)
det0_copy[..., -1] = res det0_copy[..., -1] = res
@ -275,6 +296,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)) # print("省略 :%s, lh:%s, lw:%s"%('+++' * 10, lh, lw))
# 图片的长度和宽度 # 图片的长度和宽度
imh, imw = img.shape[0:2] imh, imw = img.shape[0:2]
points = None
box = xy2xyxy(box) box = xy2xyxy(box)
# 框框左上的位置 # 框框左上的位置
x0, y1 = box[0][0], box[0][1] x0, y1 = box[0][0], box[0][1]
@ -315,22 +337,140 @@ def plot_one_box_auto(box, img, color=None, line_thickness=2, label_array=None):
return img, box return img, box
def draw_name_crowd(dets, img, color, outfontsize=20):
def draw_name_crowd(dets, img, color, label, line_thickness=2, outfontsize=20):
font = ImageFont.truetype(FONT_PATH, outfontsize, encoding='utf-8') font = ImageFont.truetype(FONT_PATH, outfontsize, encoding='utf-8')
H,W = img.shape[:2] if len(dets) == 2:
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) label = '当前人数：%d'%len(dets[0])
# img = Image.fromarray(img) detP = dets[0]
# width, height = img.size line = dets[1]
Wrate = W // 128 * 128/W for p in detP:
Hrate = H // 128 * 128/H img = cv2.circle(img, (int(p[0]), int(p[1])), line, color, -1)
label_arr = get_label_array(color, label, font, outfontsize)
lh, lw = label_arr.shape[0:2]
img[0:lh, 0:lw, :] = label_arr
elif len(dets) == 3:
detP = dets[1]
line = dets[2]
for p in detP:
img = cv2.circle(img, (int(p[0]), int(p[1])), line, color, -1)
# img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) detM = dets[0]
h, w = img.shape[:2]
for b in detM:
label = '该建筑下行人及数量：%d'%(int(b[4]))
label_arr = get_label_array(color, label, font, outfontsize)
lh, lw = label_arr.shape[0:2]
# 框框左上的位置
x0, y1 = int(b[0]), int(b[1])
# print("省略 :%s, x0:%s, y1:%s"%('+++' * 10, x0, y1))
x1, y0 = x0 + lw, y1 - lh
# 如果y0小于0, 说明超过上边框
if y0 < 0:
y0 = 0
# y1等于文字高度
y1 = y0 + lh
# 如果y1框框的高大于图片高度
if y1 > h:
# y1等于图片高度
y1 = h
# y0等于y1减去文字高度
y0 = y1 - lh
# 如果x0小于0
if x0 < 0:
x0 = 0
x1 = x0 + lw
if x1 > w:
x1 = w
x0 = x1 - lw
for p in dets: cv2.polylines(img, [np.asarray(xy2xyxy(b), np.int32)], True, (0, 128, 255), 2)
img = cv2.circle(img, (int(p[0]/Wrate), int(p[1]/Hrate)), line_thickness, color, -1) img[y0:y1, x0:x1, :] = label_arr
Calc_label_arr = get_label_array(color, label, font, outfontsize)
lh, lw = Calc_label_arr.shape[0:2]
img[0:lh, 0:lw, :] = Calc_label_arr
return img, dets
return img, dets
def draw_name_points(img,box,color):
font = ImageFont.truetype(FONT_PATH, 6, encoding='utf-8')
points = box[-1]
arrea = cv2.contourArea(points)
label = '火焰'
arealabel = '面积:%s' % f"{arrea:.1e}"
label_array_area = get_label_array(color, arealabel, font, 10)
label_array = get_label_array(color, label, font, 10)
lh_area, lw_area = label_array_area.shape[0:2]
box = box[:4]
# 框框左上的位置
x0, y1 = box[0][0], max(box[0][1] - lh_area - 3, 0)
x1, y0 = box[1][0], box[1][1]
x0_label, y0_label, x1_label, y1_label = get_label_left(x0, y1, label_array, img)
x0_area, y0_area, x1_area, y1_area = get_label_right(x1, y0, label_array_area)
img[y0_label:y1_label, x0_label:x1_label, :] = label_array
img[y0_area:y1_area, x0_area:x1_area, :] = label_array_area
# cv2.drawContours(img, points, -1, color, tl)
cv2.polylines(img, [points], False, color, 2)
if lw_area < box[1][0] - box[0][0]:
box = [(x0, y1), (x1, y1), (x1, box[2][1]), (x0, box[2][1])]
else:
box = [(x0_label, y1), (x1, y1), (x1, box[2][1]), (x0_label, box[2][1])]
box = np.asarray(box, np.int32)
cv2.polylines(img, [box], True, color, 2)
return img, box
def draw_name_border(box,color,label_array,border):
box = xywh2xyxy(box[:4])
cx, cy = int((box[0][0] + box[2][0]) / 2), int((box[0][1] + box[2][1]) / 2)
flag = cv2.pointPolygonTest(border, (int(cx), int(cy)),
False) # 若为False，会找点是否在内，外，或轮廓上
if flag == 1:
color = [0, 0, 255]
# 纯白色是(255, 255, 255)，根据容差定义白色范围
lower_white = np.array([255 - 30] * 3, dtype=np.uint8)
upper_white = np.array([255, 255, 255], dtype=np.uint8)
# 创建白色区域的掩码（白色区域为True，非白色为False）
white_mask = cv2.inRange(label_array, lower_white, upper_white)
# 创建与原图相同大小的目标颜色图像
target_img = np.full_like(label_array, color, dtype=np.uint8)
# 先将非白色区域设为目标颜色，再将白色区域覆盖回原图颜色
label_array = np.where(white_mask[..., None], label_array, target_img)
return color,label_array
def draw_transparent_red_polygon(img, points, alpha=0.5):
"""
在图像中指定的多边形区域绘制半透明红色
参数:
image_path: 原始图像路径
points: 多边形顶点坐标列表，格式为[(x1,y1), (x2,y2), ..., (xn,yn)]
output_path: 输出图像路径
alpha: 透明度系数，0-1之间，值越小透明度越高
"""
# 读取原始图像
if img is None:
raise ValueError(f"无法读取图像")
# 创建与原图大小相同的透明图层（RGBA格式）
overlay = np.zeros((img.shape[0], img.shape[1], 4), dtype=np.uint8)
# 将点列表转换为适合cv2.fillPoly的格式
#pts = np.array(points, np.int32)
pts = points.reshape((-1, 1, 2))
# 在透明图层上绘制红色多边形（BGR为0,0,255）
# 最后一个通道是Alpha值，控制透明度，黄色rgb
cv2.fillPoly(overlay, [pts], (255, 0, 0, int(alpha * 255)))
# 将透明图层转换为BGR格式（用于与原图混合）
overlay_bgr = cv2.cvtColor(overlay, cv2.COLOR_RGBA2BGR)
# 创建掩码，用于提取红色区域
mask = overlay[:, :, 3] / 255.0
mask = np.stack([mask] * 3, axis=-1) # 转换为3通道
# 混合原图和透明红色区域
img = img * (1 - mask) + overlay_bgr * mask
img = img.astype(np.uint8)
# # 保存结果
# cv2.imwrite(output_path, result)
return img