1536 lines
60 KiB
Python
1536 lines
60 KiB
Python
import asyncio
|
||
import json
|
||
|
||
import cv2
|
||
import subprocess
|
||
from threading import Thread, Lock, Event
|
||
import time
|
||
import queue
|
||
import numpy as np
|
||
|
||
import datetime
|
||
import os
|
||
from ultralytics import YOLO # 导入 Ultralytics YOLO 模型
|
||
|
||
from func_aimodellist.func_1x import func_100004, func_100000, func_100021
|
||
from middleware.minio_util import upload_file
|
||
# from Ai_tottle.minio_oss import upload_file
|
||
from mqtt_pub import MQTTClient
|
||
|
||
from ultralytics import solutions
|
||
import torch
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
|
||
from yolo.yolo_multi_model import YOLOModel, convert_result_to_multi_result, multi_model_process_frame
|
||
import av
|
||
|
||
|
||
|
||
# 全局变量
|
||
ifAI = {'status': False}
|
||
deskLock = Lock()
|
||
frame_queue = queue.Queue(maxsize=60) # 增加帧缓冲队列大小
|
||
processed_frame_queue = queue.Queue(maxsize=30) # 处理后的帧队列
|
||
|
||
list_track_id = [] # 缓存track_id 进而方便做事件判断
|
||
|
||
stop_event = Event()
|
||
|
||
mqtt_client = None
|
||
# MQTT 代理地址和端口
|
||
# broker = "112.44.103.230" # 公共 MQTT 代理（免费）
|
||
broker = "8.137.54.85" # 公共 MQTT 代理（免费）
|
||
port = 1883 # MQTT 默认端口
|
||
# 主题
|
||
topic = "thing/product/ai/events"
|
||
|
||
local_drc_message = None
|
||
CHINESE_LABELS = {0: "人", 1: "车", 2: "车", 3: "猫", 4: "购", 5: "猫", 6: "猫", 7: "猫", 8: "猫", 9: "猫", 10: "猫"}
|
||
|
||
fps = 30
|
||
|
||
# 创建 RTMP 输出容器（显式指定 format="flv"）
|
||
# container = av.open(rtmp_url, mode="w", format="flv")
|
||
|
||
# 初始不设置宽高，等待第一帧获取尺寸
|
||
stream = None
|
||
container = None
|
||
|
||
first_frame_received = False
|
||
|
||
|
||
# font_path="config/SIMSUN.TTC" 为中文配置文件路径，否则标签对中文支持不好，linux上要做特殊处理
|
||
|
||
def put_chinese_text(img, text, position, font_path="config/SIMSUN.TTC", font_size=20, color=(0, 255, 0)):
|
||
|
||
# def put_chinese_text(img, text, position, font_path="simhei.ttf", font_size=20, color=(0, 255, 0)):
|
||
"""使用PIL库在图片上绘制中文"""
|
||
img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
|
||
draw = ImageDraw.Draw(img_pil)
|
||
|
||
try:
|
||
font = ImageFont.truetype(font_path, font_size)
|
||
except:
|
||
font = ImageFont.load_default()
|
||
|
||
draw.text(position, text, font=font, fill=color)
|
||
return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
|
||
|
||
|
||
# 针对tracker的结果做统计输出
|
||
def cal_tricker_results(frame, results, class_names,
|
||
func_id, local_func_cache, para, model_cls,
|
||
chinese_label,int_func_id):
|
||
global list_track_id
|
||
"""Draw detection results on the frame."""
|
||
|
||
list_func_id=[int_func_id]
|
||
annotated_frame = frame.copy()
|
||
|
||
result_boxes = results.boxes
|
||
result_clss = results.clss
|
||
result_conf = results.confs
|
||
|
||
# print(f"result_clss {len(result_boxes)} {len(result_clss)} {len(result_conf)}")
|
||
# print("cal_t 1")
|
||
for i, box in enumerate(result_boxes):
|
||
class_id = result_clss[i]
|
||
conf = result_conf[i]
|
||
if isinstance(box, (torch.Tensor, np.ndarray)):
|
||
x1, y1, x2, y2 = map(int, box[:4])
|
||
else:
|
||
x1, y1, x2, y2 = box[0], box[1], box[2], box[3]
|
||
label = chinese_label.get(class_id, str(class_id))
|
||
text = f"{label} {conf:.2f}"
|
||
|
||
# 绘制边界框
|
||
cv2.rectangle(annotated_frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
|
||
|
||
# 计算文本位置（确保在框上方显示）
|
||
text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0]
|
||
text_width, text_height = text_size[0], text_size[1]
|
||
|
||
# 调整文本位置到框上方
|
||
text_x = x1
|
||
text_y = y1 - 5 # 稍微向上移动5像素
|
||
|
||
# 如果文本超出图像顶部，则放在框内部下方
|
||
if text_y < 0:
|
||
text_y = y2 + text_height + 5
|
||
|
||
# 使用PIL绘制中文标签（先绘制背景矩形增强可读性）
|
||
# 创建背景矩形
|
||
bg_color = (0, 0, 0) # 黑色背景
|
||
padding = 2
|
||
annotated_frame = cv2.rectangle(
|
||
annotated_frame,
|
||
(int(text_x) - padding, int(text_y) - text_height - padding),
|
||
(int(text_x) + text_width + padding, int(text_y) + padding),
|
||
bg_color,
|
||
-1
|
||
)
|
||
|
||
# 再绘制文本（使用PIL绘制中文）
|
||
temp_img = annotated_frame.copy()
|
||
annotated_frame = put_chinese_text(
|
||
temp_img,
|
||
text,
|
||
(text_x, text_y - text_height), # PIL的坐标系不同，需要调整
|
||
font_size=15,
|
||
color=(0, 255, 0)
|
||
)
|
||
|
||
box_result = []
|
||
|
||
# print("cal_t 2")
|
||
# 当前是伪代码，针对人车识别模型
|
||
# human_pt={0: 'pedestrian', 1: 'people', 2: 'bicycle', 3: 'car', 4: 'van', 5: 'truck', 6: 'tricycle', 7: 'awning-tricycle', 8: 'bus', 9: 'motor'}
|
||
# human_pt = {0: "行人", 1: "人", 2: "自行车", 3: "车", 4: "面包车", 5: "货车", 6: "三轮车", 7: "带篷三轮车",
|
||
# 8: "巴士", 9: "摩托车"}
|
||
|
||
# 接下来部分是功能筛选代码，针对100001、100002 功能代码，做部分逻辑判断
|
||
# 发现人 即报警
|
||
# 伪代码，暂时写死
|
||
if 100014 in list_func_id:
|
||
list_func_id.append(100000)
|
||
|
||
func_id_100000 = 100000
|
||
if func_id_100000 in list_func_id:
|
||
# index = [0,1]
|
||
# type_name_list=[]
|
||
# for i in model_cls:
|
||
# type_name_list.append(human_pt.get(i))
|
||
result_100001 = func_100000(results, model_cls, chinese_label, func_id_100000, list_track_id,func_id)
|
||
|
||
if result_100001 is not None:
|
||
# print(f"result_100001result_100001 {result_100001}")
|
||
box_result.append(result_100001)
|
||
# box_result["result_100001"] = result_100001
|
||
|
||
# print("cal_t 3")
|
||
func_id_100002 = 100002
|
||
if func_id_100002 in list_func_id:
|
||
# index = [2,3,4,5,6,7,8,9]
|
||
# type_name_list=[]
|
||
# for i in index:
|
||
# type_name_list.append(human_pt.get(i))
|
||
# print("c 1")
|
||
result_100002 = func_100000(results, model_cls, chinese_label, func_id_100002, list_track_id)
|
||
# print("c 2")
|
||
if result_100002 is not None:
|
||
# print(f"result_100001result_100001 {result_100002}")
|
||
box_result.append(result_100002)
|
||
# box_result["result_100001"] = result_100001
|
||
|
||
func_id_100004 = 100004
|
||
if func_id_100004 in list_func_id:
|
||
# index = [0,1]
|
||
# type_name_list=[]
|
||
local_cache = local_func_cache["func_100004"]
|
||
# for i in index:
|
||
# type_name_list.append(human_pt.get(i))
|
||
result_100004, local_cache = func_100004(results, model_cls, chinese_label, func_id_100004, list_track_id,
|
||
local_cache)
|
||
if result_100004 is not None:
|
||
local_func_cache["func_100004"] = local_cache # 更新全局缓存
|
||
box_result.append(result_100004)
|
||
# print(f"result_100004 {result_100004}")
|
||
|
||
# print("cal_t 4")
|
||
func_id_100006 = 100006
|
||
if func_id_100006 in list_func_id:
|
||
# index = [2,3,4,5,6,7,8,9]
|
||
# type_name_list=[]
|
||
local_cache = local_func_cache["func_100006"]
|
||
# for i in index:
|
||
# type_name_list.append(human_pt.get(i))
|
||
result_100006, local_cache = func_100004(results, model_cls, chinese_label, func_id_100006, list_track_id,
|
||
local_cache)
|
||
if result_100006 is not None:
|
||
local_func_cache["func_100006"] = local_cache # 更新全局缓存
|
||
box_result.append(result_100006)
|
||
# print(f"result_100006 {result_100006}")
|
||
# print("cal_t 5")
|
||
|
||
func_id_100021 = 100021
|
||
if func_id_100021 in list_func_id:
|
||
# index = [0,1]
|
||
# type_name_list=[]
|
||
number_n = para["N"] # rest接口中的N值，做人员聚集、车辆聚集的判断参数
|
||
# local_cache=local_func_cache["func_id_100021"]
|
||
# for i in index:
|
||
# type_name_list.append(human_pt.get(i))
|
||
result_100021 = func_100021(results, model_cls, chinese_label, func_id_100021, number_n)
|
||
if result_100021 is not None:
|
||
box_result.append(result_100021)
|
||
# print(f"result_100021 {result_100021}")
|
||
|
||
func_id_100023 = 100023
|
||
if func_id_100023 in list_func_id:
|
||
# index = [2,3,4,5,6,7,8,9]
|
||
# type_name_list=[]
|
||
number_n = para["N"] # rest接口中的N值，做人员聚集、车辆聚集的判断参数
|
||
# local_cache=local_func_cache["func_100006"]
|
||
# for i in index:
|
||
# type_name_list.append(human_pt.get(i))
|
||
result_100023 = func_100021(results, model_cls, chinese_label, func_id_100023, number_n)
|
||
if result_100023 is not None:
|
||
# local_func_cache["func_100006"]=local_cache #更新全局缓存
|
||
box_result.append(result_100023)
|
||
# print(f"result_100023 {result_100023}")
|
||
|
||
# bridge_pt={0:"蜂窝",1:"剥落",2:"空腔",3:"锈蚀",4:"钢筋裸露",5:"裂缝"}
|
||
|
||
func_id_100031 = 100031
|
||
if func_id_100031 in list_func_id:
|
||
# index = [0,1,2,3,4,5]
|
||
# type_name_list=[]
|
||
# for i in index:
|
||
# type_name_list.append(bridge_pt.get(i))
|
||
result_100031 = func_100000(results, model_cls, chinese_label, func_id_100031, list_track_id)
|
||
if result_100031 is not None:
|
||
# print(f"result_100031result_100031 {result_100031}")
|
||
box_result.append(result_100031)
|
||
|
||
# smoke_pt={0:"烟雾"}
|
||
|
||
func_id_100041 = 100041
|
||
if func_id_100041 in list_func_id:
|
||
# index = [0]
|
||
# type_name_list=[]
|
||
# for i in index:
|
||
# type_name_list.append(smoke_pt.get(i))
|
||
result_100041 = func_100000(results, model_cls, chinese_label, func_id_100041, list_track_id)
|
||
if result_100041 is not None:
|
||
# print(f"result_100041result_100041 {result_100041}")
|
||
box_result.append(result_100041)
|
||
|
||
# hwrc_pt={0:"人",1:"卡车",2:"汽车",3:"自行车"}
|
||
|
||
func_id_100061 = 100061
|
||
if func_id_100061 in list_func_id:
|
||
# index = [0,1,2,3]
|
||
# type_name_list=[]
|
||
# for i in index:
|
||
# type_name_list.append(hwrc_pt.get(i))
|
||
result_100061 = func_100000(results, model_cls, chinese_label, func_id_100061, list_track_id)
|
||
if result_100061 is not None:
|
||
# print(f"result_100061result_100061 {result_100061}")
|
||
box_result.append(result_100061)
|
||
|
||
# hwgf_pt={0:"单一热班",1:"大面积热斑",2:"单一热班&异常低温",3:"大面积热班&异常低温",4:"异常低温",5:"单一热斑&遮挡",6:"异常低温&热斑",7:"二极管短路",8:"阳光反射",9:"二极管短路&异常低温"}
|
||
|
||
func_id_100081 = 100081
|
||
if func_id_100081 in list_func_id:
|
||
# index = [0,1,2,3,4,5,6,7,8,9]
|
||
# type_name_list=[]
|
||
# for i in index:
|
||
# type_name_list.append(hwgf_pt.get(i))
|
||
result_100081 = func_100000(results, model_cls, chinese_label, func_id_100081, list_track_id)
|
||
if result_100081 is not None:
|
||
# print(f"result_100081result_100081 {result_100081}")
|
||
box_result.append(result_100081)
|
||
|
||
# # # 缩放一下，好展示
|
||
# resized_img = cv2.resize(annotated_frame, (0, 0), fx=0.5, fy=0.5) # 按比例缩放
|
||
# cv2.imshow("annotated_frame",resized_img)
|
||
# cv2.waitKey(0) # 等待任意按键，参数为延迟毫秒（0表示无限等待）
|
||
# cv2.destroyAllWindows() # 关闭窗口
|
||
|
||
|
||
func_id_100008 = 100008
|
||
if func_id_100008 in list_func_id:
|
||
# index = [0,1]
|
||
# type_name_list=[]
|
||
# for i in model_cls:
|
||
# type_name_list.append(human_pt.get(i))
|
||
result_100008 = func_100000(results, model_cls, chinese_label, func_id_100008, list_track_id)
|
||
if result_100008 is not None:
|
||
# print(f"result_100001result_100001 {result_100001}")
|
||
box_result.append(result_100008)
|
||
|
||
|
||
func_id_100052 = 100052
|
||
if func_id_100052 in list_func_id:
|
||
# index = [0,1]
|
||
# type_name_list=[]
|
||
# for i in model_cls:
|
||
# type_name_list.append(human_pt.get(i))
|
||
result_100052 = func_100000(results, model_cls, chinese_label, func_id_100052, list_track_id)
|
||
if result_100052 is not None:
|
||
# print(f"result_100001result_100001 {result_100001}")
|
||
box_result.append(result_100052)
|
||
|
||
func_id_100091 = 100091
|
||
if func_id_100091 in list_func_id:
|
||
# index = [0,1]
|
||
# type_name_list=[]
|
||
# for i in model_cls:
|
||
# type_name_list.append(human_pt.get(i))
|
||
result_100091 = func_100000(results, model_cls, chinese_label, func_id_100091, list_track_id)
|
||
if result_100091 is not None:
|
||
# print(f"result_100001result_100001 {result_100001}")
|
||
box_result.append(result_100091)
|
||
|
||
func_id_100092 = 100092
|
||
if func_id_100092 in list_func_id:
|
||
# index = [0,1]
|
||
# type_name_list=[]
|
||
# for i in model_cls:
|
||
# type_name_list.append(human_pt.get(i))
|
||
result_100092 = func_100000(results, model_cls, chinese_label, func_id_100092, list_track_id)
|
||
if result_100092 is not None:
|
||
# print(f"result_100001result_100001 {result_100001}")
|
||
box_result.append(result_100092)
|
||
|
||
|
||
|
||
return annotated_frame, box_result
|
||
|
||
|
||
# 输入为检测结果，输出为检测分类，统计后方便前端使用
|
||
def extract_box_details(last_results, model_id):
|
||
"""
|
||
从 YOLO 检测结果中提取边界框的详细信息
|
||
:param last_results: YOLO 检测结果对象
|
||
:param model_id: 模型id
|
||
:return: 包含边界框事件和数量的字典
|
||
"""
|
||
# box_event = []
|
||
box_count = []
|
||
count_dict = {}
|
||
human_pt = {0: "行人", 1: "人", 2: "自行车", 3: "车", 4: "面包车", 5: "货车", 6: "三轮车", 7: "带篷三轮车",
|
||
8: "巴士", 9: "摩托车"}
|
||
for box in last_results.boxes:
|
||
cls_index = int(box.cls[0]) # 类别索引
|
||
if cls_index in count_dict:
|
||
count_dict[cls_index] += 1
|
||
else:
|
||
count_dict[cls_index] = 1
|
||
|
||
for key, value in count_dict.items():
|
||
type_name = human_pt.get(key)
|
||
box_count.append({
|
||
"type": key,
|
||
"type_name": type_name, # 伪代码，后续需要修改
|
||
"count": value
|
||
})
|
||
|
||
box_detail = {
|
||
"model_id": model_id,
|
||
"box_count": box_count
|
||
}
|
||
|
||
return box_detail
|
||
|
||
|
||
def setIfAI(pb1):
|
||
deskLock.acquire()
|
||
ifAI['status'] = pb1
|
||
deskLock.release()
|
||
|
||
|
||
def getIfAI():
|
||
return ifAI['status']
|
||
|
||
|
||
def stopAIVideo():
|
||
print("正在停止AI视频处理...")
|
||
setIfAI(False)
|
||
stop_event.set()
|
||
|
||
# 等待足够长的时间确保资源释放
|
||
wait_count = 0
|
||
max_wait = 5 # 减少最大等待时间到5秒
|
||
|
||
while stop_event.is_set() and wait_count < max_wait:
|
||
time.sleep(0.5)
|
||
wait_count += 1
|
||
|
||
if wait_count >= max_wait:
|
||
print("警告: 停止AI视频处理超时，强制终止")
|
||
# 不使用_thread._interrupt_main()，改用其他方式强制终止
|
||
try:
|
||
# 尝试终止可能运行的进程
|
||
import os
|
||
import signal
|
||
import psutil
|
||
|
||
# 查找并终止可能的FFmpeg进程
|
||
current_process = psutil.Process(os.getpid())
|
||
for child in current_process.children(recursive=True):
|
||
try:
|
||
child_name = child.name().lower()
|
||
if 'ffmpeg' in child_name:
|
||
print(f"正在终止子进程: {child.pid} ({child_name})")
|
||
child.send_signal(signal.SIGTERM)
|
||
except:
|
||
pass
|
||
except:
|
||
pass
|
||
else:
|
||
print("AI视频处理已停止")
|
||
|
||
|
||
def startBackAIVideo(task_id, video_path,push_url, m1, model_cls, chinese_label, list_func_id, confidence, para, mqtt_ip,
|
||
mqtt_port, mqtt_topic):
|
||
if ifAI['status']:
|
||
stopAIVideo()
|
||
time.sleep(1)
|
||
stop_event.clear()
|
||
thread = Thread(target=startBackAIVideo2,
|
||
args=(
|
||
task_id, video_path, push_url,m1, list_func_id, confidence, para, mqtt_ip, mqtt_port, mqtt_topic,
|
||
model_cls,
|
||
chinese_label))
|
||
# cls2_thread = Thread(target=cls2_find, args=(video_path,m1, cls, confidence))
|
||
# cls2_thread.daemon = True # 守护线程，主程序退出时线程也会退出
|
||
thread.daemon = True # 守护线程，主程序退出时线程也会退出
|
||
|
||
thread.start()
|
||
|
||
|
||
def read_frames(cap,push_url, frame_queue):
|
||
"""优化的帧读取线程"""
|
||
global first_frame_received,stream,container
|
||
if container is None:
|
||
container = av.open(push_url, mode="w", format="flv")
|
||
frame_count = 0
|
||
last_time = time.time()
|
||
last_fps_time = time.time()
|
||
# 减小目标帧间隔时间，提高读取帧率
|
||
target_time_per_frame = 1.0 / 60.0 # 目标帧间隔时间（提高到60fps）
|
||
|
||
# 添加连接断开检测
|
||
connection_error_count = 0
|
||
max_connection_errors = 10 # 最多允许连续10次连接错误
|
||
last_successful_read = time.time()
|
||
max_read_wait = 60.0 # 30秒无法读取则认为连接断开
|
||
|
||
# 预先丢弃几帧，确保从新帧开始处理
|
||
for _ in range(5):
|
||
cap.grab()
|
||
retry_count = 0
|
||
|
||
max_retries = 3
|
||
|
||
while not stop_event.is_set() and retry_count < max_retries:
|
||
try:
|
||
# 确保container使用正确的push_url
|
||
if container is None:
|
||
try:
|
||
container = av.open(push_url, mode="w", format="flv")
|
||
print(f"成功打开RTMP推流: {push_url}")
|
||
except Exception as e:
|
||
print(f"无法打开RTMP推流地址: {push_url}, 错误: {e}")
|
||
stop_event.set()
|
||
return
|
||
|
||
current_time = time.time()
|
||
elapsed_time = current_time - last_time
|
||
# 检查是否长时间无法读取帧
|
||
if current_time - last_successful_read > max_read_wait:
|
||
print(f"超时警告: {max_read_wait}秒内未能读取到有效帧，可能连接已断开")
|
||
stop_event.set()
|
||
break
|
||
# 帧率控制，但更积极地读取
|
||
if elapsed_time < target_time_per_frame:
|
||
time.sleep(target_time_per_frame - elapsed_time)
|
||
continue
|
||
# print(f"4 {frame_queue.qsize()} {frame_queue.maxsize * 0.8}")
|
||
# 当队列快满时，跳过一些帧以避免延迟累积
|
||
if frame_queue.qsize() > frame_queue.maxsize * 0.8:
|
||
# 跳过一些帧
|
||
cap.grab()
|
||
last_time = time.time()
|
||
|
||
# 清空队列
|
||
while not frame_queue.empty():
|
||
try:
|
||
frame_queue.get_nowait() # 非阻塞地获取并移除队列中的元素
|
||
except queue.Empty:
|
||
break # 防止在极端情况下出现的竞争条件
|
||
# print(f"5 {frame_queue.qsize()} {frame_queue.maxsize * 0.8}")
|
||
# print(f"引入solutions.ObjectCounter，单帧处理速度变慢，队列堆积后，直接清空")
|
||
continue
|
||
# print(6)
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
print("拉流错误：无法读取帧")
|
||
connection_error_count += 1
|
||
if connection_error_count >= max_connection_errors:
|
||
print(f"连续{max_connection_errors}次无法读取帧，可能连接已断开，正在停止流程...")
|
||
stop_event.set()
|
||
break
|
||
time.sleep(0.5) # 短暂等待后重试
|
||
continue
|
||
|
||
# 成功读取了帧，重置错误计数
|
||
connection_error_count = 0
|
||
last_successful_read = time.time()
|
||
|
||
frame_count += 1
|
||
if frame_count % 60 == 0: # 每60帧计算一次FPS
|
||
current_fps_time = time.time()
|
||
fps = 60 / (current_fps_time - last_fps_time)
|
||
# print(f"拉流FPS: {fps:.2f}")
|
||
last_fps_time = current_fps_time
|
||
|
||
last_time = time.time()
|
||
frame_queue.put((frame, time.time())) # 添加时间戳
|
||
height, width, channels = frame.shape
|
||
|
||
# 如果是第一帧，初始化流参数
|
||
if not first_frame_received:
|
||
stream = container.add_stream("libx264", rate=30)
|
||
stream.width = width
|
||
stream.height = height
|
||
stream.pix_fmt = "yuv420p"
|
||
stream.options = {
|
||
"tune": "zerolatency",
|
||
"preset": "ultrafast",
|
||
"crf": "23",
|
||
}
|
||
first_frame_received = True
|
||
print(f"Initialized stream with resolution: {width}x{height}")
|
||
|
||
if first_frame_received:
|
||
# 将 NumPy 数组转换为 AV 帧
|
||
av_frame = av.VideoFrame.from_ndarray(frame, format="bgr24")
|
||
|
||
# 编码并推送帧
|
||
for packet in stream.encode(av_frame):
|
||
container.mux(packet)
|
||
|
||
|
||
except av.error.BrokenPipeError:
|
||
print("RTMP 连接断开，尝试重连...")
|
||
retry_count += 1
|
||
if container:
|
||
container.close()
|
||
container = None
|
||
time.sleep(2) # 等待2秒后重试
|
||
except Exception as e:
|
||
print(f"读取帧时发生错误: {e}")
|
||
break
|
||
|
||
|
||
|
||
|
||
|
||
def process_frames_tricker(push_url,frame_queue, processed_frame_queue, model_path, list_func_id, confidence, para, model_cls,
|
||
chinese_label, use_fp16=False):
|
||
"""Count specific classes of objects in a video with dynamic processing adjustments."""
|
||
global first_frame_received,stream,container
|
||
if container is None:
|
||
container = av.open(push_url, mode="w", format="flv")
|
||
|
||
# solutions默认会在图上绘制一个方框，很烦
|
||
line_points = [(-10, -10), (-10, -10)]
|
||
# 伪代码，后续识别类型需要跟ai_model_list 的id关联
|
||
# cls = [0, 1, 2, 3, 4]
|
||
# 这里参数show=True，一旦通过rest接口重启程序，会导致线程锁住
|
||
# counter = solutions.ObjectCounter(show=False, region=line_points, model=model_path, classes=cls)
|
||
counter = solutions.ObjectCounter(show=False, region=line_points, model=model_path, classes=model_cls)
|
||
|
||
|
||
# Model processing setup
|
||
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
|
||
try:
|
||
counter.model.to(device)
|
||
if hasattr(counter.model, 'args') and hasattr(counter.model.args, 'batch'):
|
||
counter.model.args.batch = 1
|
||
# 启用半精度
|
||
if use_fp16 and device.startswith('cuda') and hasattr(counter.model, 'model'):
|
||
try:
|
||
counter.model.model = counter.model.model.half()
|
||
print("✅ 启用半精度 FP16 模式")
|
||
except Exception as half_err:
|
||
print(f"⚠️ 半精度转换失败: {half_err}")
|
||
else:
|
||
print("ℹ️ 半精度模式未启用")
|
||
|
||
except Exception as e:
|
||
print(f"⚠️ 模型设备配置警告: {e}")
|
||
|
||
frame_count = 0
|
||
process_times = []
|
||
last_results = None
|
||
last_counter = None
|
||
skip_counter = 0
|
||
max_skip = 2
|
||
error_count = 0
|
||
|
||
# Get class names from the model
|
||
class_names = counter.model.names if hasattr(counter.model, 'names') else {i: str(i) for i in range(1000)}
|
||
|
||
local_func_cache = {
|
||
"func_100000": None,
|
||
"func_100004": None, # 存储缓存，缓存人员track_id
|
||
"func_100006": None # 存储缓存，缓存车辆track_id
|
||
|
||
}
|
||
|
||
while not stop_event.is_set():
|
||
try:
|
||
if processed_frame_queue.qsize() >= processed_frame_queue.maxsize * 0.8:
|
||
time.sleep(0.01)
|
||
continue
|
||
|
||
frame, timestamp = frame_queue.get(timeout=0.2)
|
||
frame_copy = frame.copy()
|
||
if time.time() - timestamp > 0.3:
|
||
continue
|
||
|
||
frame_count += 1
|
||
|
||
if skip_counter > 0 and last_counter is not None:
|
||
skip_counter -= 1
|
||
annotated_frame, box_result = cal_tricker_results(frame, counter, class_names, list_func_id,
|
||
local_func_cache, para, model_cls, chinese_label)
|
||
continue
|
||
|
||
process_start = time.time()
|
||
|
||
resize_scale = 1.0
|
||
qsize = frame_queue.qsize()
|
||
maxsize = frame_queue.maxsize
|
||
|
||
if qsize > maxsize * 0.7:
|
||
resize_scale = 0.4
|
||
elif qsize > maxsize * 0.5:
|
||
resize_scale = 0.6
|
||
elif qsize > maxsize * 0.3:
|
||
resize_scale = 0.8
|
||
|
||
if resize_scale < 1.0:
|
||
process_frame = cv2.resize(frame, (0, 0), fx=resize_scale, fy=resize_scale)
|
||
else:
|
||
process_frame = frame
|
||
|
||
if process_frame is None or process_frame.size == 0:
|
||
print("警告: 输入帧无效")
|
||
continue
|
||
|
||
try:
|
||
# Run object counting/tracking
|
||
# 获取当前时间的时间戳，包含毫秒
|
||
# timestamp = time.time()
|
||
# # 获取13位长度的时间戳（毫秒级）
|
||
# time_start = int(timestamp * 1000)
|
||
# frame_copy=process_frame # 暂存frame，方便进行结果重绘
|
||
timestart = time.time()
|
||
results = counter(process_frame)
|
||
timeend = time.time()
|
||
print(f"timestart-timeend {timeend-timestart}")
|
||
if results is None: # 处理返回None的情况
|
||
print("警告: 检测结果为None，使用上一帧结果")
|
||
# results = getattr(counter, 'last_results', None)
|
||
# if results is None:
|
||
continue
|
||
# timestamp = time.time()
|
||
# # 获取13位长度的时间戳（毫秒级）
|
||
# time_end = int(timestamp * 1000)
|
||
# print(f"time_end - time_start {time_end - time_start}")
|
||
# print("1")
|
||
last_counter = counter # 基于计数的方法，识别后的结果存储在counter中
|
||
# print("2")
|
||
annotated_frame, box_result = cal_tricker_results(
|
||
frame_copy if resize_scale == 1.0 else cv2.resize(frame_copy, (frame.shape[1], frame.shape[0])),
|
||
counter, class_names, list_func_id, local_func_cache, para, model_cls, chinese_label)
|
||
|
||
# if annotated_frame is None:
|
||
# print("annotated_frame is None")
|
||
#
|
||
# if box_result is None:
|
||
# print("box_result is None")
|
||
|
||
# print("3")
|
||
processed_frame_queue.put((annotated_frame, timestamp, box_result))
|
||
# print_box_details(results)
|
||
# print("=4")
|
||
|
||
|
||
|
||
|
||
if qsize > maxsize * 0.5:
|
||
skip_counter = max_skip
|
||
|
||
except Exception as infer_err:
|
||
print(f"推理错误1: {infer_err}")
|
||
|
||
if last_counter is not None:
|
||
annotated_frame, box_result = cal_tricker_results(frame, counter, class_names, list_func_id,
|
||
local_func_cache, para, model_cls, chinese_label)
|
||
else:
|
||
annotated_frame = frame.copy()
|
||
|
||
process_end = time.time()
|
||
process_times.append(process_end - process_start)
|
||
if len(process_times) > 30:
|
||
process_times.pop(0)
|
||
|
||
if frame_count % 30 == 0:
|
||
avg_process_time = sum(process_times) / len(process_times)
|
||
fps = 1.0 / avg_process_time if avg_process_time > 0 else 0
|
||
print(
|
||
f"模型处理FPS: {fps:.2f}, 平均处理时间: {avg_process_time * 1000:.2f}ms, 队列大小: {qsize}, 缩放比例: {resize_scale:.2f}")
|
||
|
||
processed_frame_queue.put((annotated_frame, timestamp, box_result))
|
||
error_count = 0
|
||
|
||
except queue.Empty:
|
||
continue
|
||
except Exception as e:
|
||
print(f"处理帧错误: {e}")
|
||
error_count += 1
|
||
if error_count >= 5:
|
||
print(f"连续处理错误达到5次 ，正在停止处理...")
|
||
break
|
||
# finally:
|
||
# # 清理资源
|
||
# for packet in stream.encode():
|
||
# container.mux(packet)
|
||
# container.close()
|
||
|
||
|
||
def multi_model_process_frames(frame_queue, processed_frame_queue, model_path, list_func_id, confidence, para,
|
||
model_cls, chinese_label, use_fp16=False):
|
||
"""Count specific classes of objects in a video with dynamic processing adjustments."""
|
||
models = [
|
||
YOLOModel('pt/best.pt', cls_map={'pedestrian': '行人', 'pedestrian': '行人'},
|
||
allowed_classes=['pedestrian', 'people'])
|
||
,
|
||
YOLOModel('pt/best.pt', cls_map={'car': '车'}, allowed_classes=['car'])
|
||
]
|
||
|
||
# 用于存储处理结果的字典
|
||
result_dict = {}
|
||
frame_count = 0
|
||
processing_times = []
|
||
|
||
# 启动所有模型的工作线程
|
||
for model in models:
|
||
model.start()
|
||
|
||
local_func_cache = {
|
||
"func_100000": None,
|
||
"func_100004": None, # 存储缓存，缓存人员track_id
|
||
"func_100006": None # 存储缓存，缓存车辆track_id
|
||
|
||
}
|
||
|
||
max_skip = 2
|
||
error_count = 0
|
||
while not stop_event.is_set():
|
||
try:
|
||
if processed_frame_queue.qsize() >= processed_frame_queue.maxsize * 0.8:
|
||
time.sleep(0.01)
|
||
continue
|
||
|
||
frame, timestamp = frame_queue.get(timeout=0.2)
|
||
frame_copy = frame.copy()
|
||
if time.time() - timestamp > 0.3:
|
||
continue
|
||
|
||
process_start = time.time()
|
||
|
||
resize_scale = 1.0
|
||
qsize = frame_queue.qsize()
|
||
maxsize = frame_queue.maxsize
|
||
|
||
# if qsize > maxsize * 0.7:
|
||
# resize_scale = 0.4
|
||
# elif qsize > maxsize * 0.5:
|
||
# resize_scale = 0.6
|
||
# elif qsize > maxsize * 0.3:
|
||
# resize_scale = 0.8
|
||
#
|
||
# if resize_scale < 1.0:
|
||
# process_frame = cv2.resize(frame, (0, 0), fx=resize_scale, fy=resize_scale)
|
||
# else:
|
||
# process_frame = frame
|
||
|
||
# if process_frame is None or process_frame.size == 0:
|
||
# print("警告: 输入帧无效")
|
||
# continue
|
||
|
||
# 初始化多个YOLO模型
|
||
|
||
try:
|
||
timestart = time.time()
|
||
# 处理帧（多模型并行）
|
||
multi_model_process_frame(frame, models, result_dict)
|
||
timeend1 = time.time()
|
||
print(f"timeend1 {timeend1 - timestart}")
|
||
# 获取处理结果
|
||
frame_id = id(frame)
|
||
if frame_id in result_dict:
|
||
results = result_dict.pop(frame_id)
|
||
|
||
multi_result = convert_result_to_multi_result(results)
|
||
timeend2 = time.time()
|
||
print(f"timeend2 {timeend2 - timestart}")
|
||
annotated_frame, box_result = cal_tricker_results(frame, multi_result, None, list_func_id,local_func_cache, para, model_cls, chinese_label)
|
||
timeend3 = time.time()
|
||
print(f"timeend3 {timeend3 - timestart}")
|
||
time.sleep(1)
|
||
timeend4 = time.time()
|
||
print(f"timeend4 {timeend4 - timeend3}")
|
||
# cv2.imshow('annotated_frame Frame', frame)
|
||
# cv2.imwrite('output.jpg', frame)
|
||
# #
|
||
# # results = counter(process_frame)
|
||
# # if results is None: # 处理返回None的情况
|
||
# # print("警告: 检测结果为None，使用上一帧结果")
|
||
# # # results = getattr(counter, 'last_results', None)
|
||
# # # if results is None:
|
||
# # continue
|
||
# # timestamp = time.time()
|
||
# # # 获取13位长度的时间戳（毫秒级）
|
||
# # time_end = int(timestamp * 1000)
|
||
# # print(f"time_end - time_start {time_end - time_start}")
|
||
# # print("1")
|
||
# last_counter = counter # 基于计数的方法，识别后的结果存储在counter中
|
||
# # print("2")
|
||
# annotated_frame, box_result = cal_tricker_results(
|
||
# frame_copy if resize_scale == 1.0 else cv2.resize(frame_copy, (frame.shape[1], frame.shape[0])),
|
||
# counter, class_names, list_func_id,local_func_cache,para,model_cls,chinese_label)
|
||
#
|
||
# # if annotated_frame is None:
|
||
# # print("annotated_frame is None")
|
||
# #
|
||
# # if box_result is None:
|
||
# # print("box_result is None")
|
||
|
||
# print("3")
|
||
# processed_frame_queue.put((annotated_frame, timestamp, box_result))
|
||
# print_box_details(results)
|
||
# print("=4")
|
||
if qsize > maxsize * 0.5:
|
||
skip_counter = max_skip
|
||
|
||
except Exception as infer_err:
|
||
print(f"推理错误1: {infer_err}")
|
||
|
||
# if last_counter is not None:
|
||
# annotated_frame, box_result = cal_tricker_results(frame, counter, class_names, list_func_id,local_func_cache,para,model_cls,chinese_label)
|
||
# else:
|
||
# annotated_frame = frame.copy()
|
||
|
||
# process_end = time.time()
|
||
# process_times.append(process_end - process_start)
|
||
# if len(process_times) > 30:
|
||
# process_times.pop(0)
|
||
|
||
# if frame_count % 30 == 0:
|
||
# avg_process_time = sum(process_times) / len(process_times)
|
||
# fps = 1.0 / avg_process_time if avg_process_time > 0 else 0
|
||
# print(
|
||
# f"模型处理FPS: {fps:.2f}, 平均处理时间: {avg_process_time * 1000:.2f}ms, 队列大小: {qsize}, 缩放比例: {resize_scale:.2f}")
|
||
#
|
||
# processed_frame_queue.put((annotated_frame, timestamp, box_result))
|
||
error_count = 0
|
||
|
||
except queue.Empty:
|
||
continue
|
||
except Exception as e:
|
||
print(f"处理帧错误: {e}")
|
||
error_count += 1
|
||
if error_count >= 5:
|
||
print(f"连续处理错误达到5次 ，正在停止处理...")
|
||
break
|
||
|
||
|
||
#
|
||
# def process_frames(model_id, frame_queue, processed_frame_queue, ov_model, cls, confidence, use_fp16=False):
|
||
# """处理帧的线程，添加帧率控制，支持半精度FP16"""
|
||
#
|
||
# import torch
|
||
# import time
|
||
# import queue
|
||
# import cv2
|
||
#
|
||
# error_count = 0 # 添加错误计数器
|
||
# max_errors = 5 # 最大容许错误次数
|
||
# frame_count = 0
|
||
# process_times = [] # 用于计算平均处理时间
|
||
#
|
||
# # 设置YOLO模型配置，提高性能
|
||
# ov_model.conf = confidence # 设置置信度阈值
|
||
#
|
||
# # 将模型移到设备（GPU或CPU）
|
||
# device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
|
||
# try:
|
||
# ov_model.to(device)
|
||
# # 调整批处理大小为1，减少内存占用
|
||
# if hasattr(ov_model, 'args') and hasattr(ov_model.args, 'batch'):
|
||
# ov_model.args.batch = 1
|
||
# # 启用半精度
|
||
# if use_fp16 and device.startswith('cuda') and hasattr(ov_model, 'model'):
|
||
# try:
|
||
# ov_model.model = ov_model.model.half()
|
||
# print("✅ 启用半精度 FP16 模式")
|
||
# except Exception as half_err:
|
||
# print(f"⚠️ 半精度转换失败: {half_err}")
|
||
# else:
|
||
# print("ℹ️ 半精度模式未启用")
|
||
# except Exception as e:
|
||
# print(f"⚠️ 模型设备配置警告: {e}")
|
||
#
|
||
# # 缓存先前的检测结果，用于提高稳定性
|
||
# last_results = None
|
||
# skip_counter = 0
|
||
# max_skip = 2 # 最多跳过几帧不处理
|
||
#
|
||
# while not stop_event.is_set():
|
||
# try:
|
||
# if processed_frame_queue.qsize() >= processed_frame_queue.maxsize * 0.8:
|
||
# time.sleep(0.01)
|
||
# continue
|
||
#
|
||
# frame, timestamp = frame_queue.get(timeout=0.2)
|
||
#
|
||
# if time.time() - timestamp > 0.3:
|
||
# continue
|
||
#
|
||
# frame_count += 1
|
||
#
|
||
# if skip_counter > 0 and last_results is not None:
|
||
# skip_counter -= 1
|
||
# annotated_frame = last_results.plot(conf=False, line_width=1, font_size=1.5)
|
||
#
|
||
# box_detail = extract_box_details(last_results, model_id)
|
||
#
|
||
# processed_frame_queue.put((annotated_frame, timestamp, box_detail))
|
||
# continue
|
||
#
|
||
# process_start = time.time()
|
||
#
|
||
# resize_scale = 1.0
|
||
# qsize = frame_queue.qsize()
|
||
# maxsize = frame_queue.maxsize
|
||
#
|
||
# if qsize > maxsize * 0.7:
|
||
# resize_scale = 0.4
|
||
# elif qsize > maxsize * 0.5:
|
||
# resize_scale = 0.6
|
||
# elif qsize > maxsize * 0.3:
|
||
# resize_scale = 0.8
|
||
#
|
||
# if resize_scale < 1.0:
|
||
# process_frame = cv2.resize(frame, (0, 0), fx=resize_scale, fy=resize_scale)
|
||
# else:
|
||
# process_frame = frame
|
||
#
|
||
# try:
|
||
# results = ov_model(process_frame, classes=cls, show=False)
|
||
# last_results = results[0]
|
||
#
|
||
# if resize_scale < 1.0:
|
||
# annotated_frame = cv2.resize(results[0].plot(conf=False, line_width=1, font_size=1.5),
|
||
# (frame.shape[1], frame.shape[0]))
|
||
# # print("results[0].boxes.cls ", results[0].boxes.cls)
|
||
#
|
||
# else:
|
||
# annotated_frame = results[0].plot(conf=False, line_width=1, font_size=1.5)
|
||
#
|
||
# box_detail = extract_box_details(last_results, model_id)
|
||
#
|
||
# if qsize > maxsize * 0.5:
|
||
# skip_counter = max_skip
|
||
# except Exception as infer_err:
|
||
# print(f"推理错误: {infer_err}")
|
||
# box_detail = None
|
||
# if last_results is not None:
|
||
# annotated_frame = last_results.plot(conf=False, line_width=1, font_size=1.5)
|
||
# else:
|
||
# annotated_frame = frame.copy()
|
||
#
|
||
# process_end = time.time()
|
||
# process_times.append(process_end - process_start)
|
||
# if len(process_times) > 30:
|
||
# process_times.pop(0)
|
||
#
|
||
# if frame_count % 30 == 0:
|
||
# avg_process_time = sum(process_times) / len(process_times)
|
||
# fps = 1.0 / avg_process_time if avg_process_time > 0 else 0
|
||
# print(
|
||
# f"模型处理FPS: {fps:.2f}, 平均处理时间: {avg_process_time * 1000:.2f}ms, 队列大小: {qsize}, 缩放比例: {resize_scale:.2f}")
|
||
#
|
||
# processed_frame_queue.put((annotated_frame, timestamp, box_detail))
|
||
# error_count = 0
|
||
# except queue.Empty:
|
||
# continue
|
||
# except Exception as e:
|
||
# error_count += 1
|
||
# print(f"处理帧错误: {e}")
|
||
# if error_count >= max_errors:
|
||
# print(f"连续处理错误达到{max_errors}次，正在停止处理...")
|
||
# stop_event.set()
|
||
# break
|
||
|
||
|
||
# # 读取mqtt的drc消息，并且在本地维持一个最新消息
|
||
# def read_drc_mqtt(mqtt_client):
|
||
# global local_drc_message
|
||
# while not stop_event.is_set():
|
||
# time.sleep(0.05)
|
||
# # 更新全局缓存，如果要同时适配多个飞机，这里要改为队列
|
||
# mess = mqtt_client.get_messages(timeout=0.5)
|
||
# if mess is not None:
|
||
# # local_drc_message = mess
|
||
# local_drc_message = json.loads(mess) # 解析为字典
|
||
|
||
# def read_drc_mqtt(mqtt_client):
|
||
# global local_drc_message
|
||
# while not stop_event.is_set():
|
||
# time.sleep(0.05)
|
||
# try:
|
||
# # 确保 mess 是字符串
|
||
# mess = mqtt_client.get_messages(timeout=0.5)
|
||
# if mess is not None:
|
||
# try:
|
||
# local_drc_message = json.loads(mess) # 确保 mess 是字符串
|
||
# except json.JSONDecodeError as e:
|
||
# print(f"JSON 解析错误: {e}")
|
||
# except Exception as e:
|
||
# print(f"读取 DRC 消息错误: {e}")
|
||
|
||
|
||
|
||
def get_local_drc_message():
|
||
global local_drc_message
|
||
if local_drc_message is not None:
|
||
mess =local_drc_message
|
||
local_drc_message=None
|
||
return mess
|
||
return None
|
||
|
||
from asyncio_mqtt import Client
|
||
|
||
async def async_read_drc_mqtt(mqtt_ip, mqtt_port, mqtt_topic):
|
||
global local_drc_message
|
||
async with Client(mqtt_ip, mqtt_port) as client:
|
||
async with client.messages() as messages:
|
||
await client.subscribe(mqtt_topic)
|
||
async for message in messages:
|
||
if message.topic.matches(mqtt_topic):
|
||
try:
|
||
data = json.loads(message.payload)
|
||
if (data and
|
||
isinstance(data.get("data"), dict) and
|
||
"attitude_head" in data["data"]):
|
||
local_drc_message = data
|
||
except Exception as e:
|
||
logger.error(f"Error processing MQTT message: {e}")
|
||
|
||
def read_drc_mqtt(mqtt_client):
|
||
global local_drc_message
|
||
while not stop_event.is_set():
|
||
time.sleep(0.05)
|
||
# 更新全局缓存，如果要同时适配多个飞机，这里要改为队列
|
||
mess = mqtt_client.get_messages(timeout=0.5)
|
||
if mess is not None:
|
||
local_drc_message = mess
|
||
|
||
|
||
def write_frames(processed_frame_queue, pipe, size, mqtt_client, task_id):
|
||
# # 测试用做，输出识别后的图片到窗口
|
||
# cv2.namedWindow("RTMP Stream with YOLO", cv2.WINDOW_NORMAL)
|
||
# cv2.resizeWindow("RTMP Stream with YOLO", 900, 600)
|
||
global local_drc_message
|
||
"""写入帧的线程，添加平滑处理"""
|
||
last_write_time = time.time()
|
||
target_time_per_frame = 1.0 / 30.0 # 30fps
|
||
pipe_error_count = 0 # 添加错误计数
|
||
max_pipe_errors = 3 # 最大容忍错误数
|
||
frame_count = 0
|
||
last_fps_time = time.time()
|
||
skipped_frames = 0
|
||
|
||
# 使用队列存储最近几帧，用于在需要时进行插值
|
||
recent_frames = []
|
||
max_recent_frames = 5 # 增加缓存帧数量，提高平滑性
|
||
|
||
# 使用双缓冲机制提高写入速度
|
||
buffer1 = bytearray(size[0] * size[1] * 3)
|
||
buffer2 = bytearray(size[0] * size[1] * 3)
|
||
current_buffer = buffer1
|
||
|
||
# 帧率控制参数
|
||
min_frame_interval = target_time_per_frame * 0.5 # 允许的最小帧间隔
|
||
max_frame_interval = target_time_per_frame * 2.0 # 允许的最大帧间隔
|
||
|
||
while not stop_event.is_set():
|
||
try:
|
||
# 获取处理后的帧，超时时间较短以便更平滑地处理
|
||
frame, timestamp, box_detail = processed_frame_queue.get(timeout=0.05)
|
||
|
||
# 存储最近的帧用于插值
|
||
recent_frames.append(frame)
|
||
if len(recent_frames) > max_recent_frames:
|
||
recent_frames.pop(0)
|
||
|
||
current_time = time.time()
|
||
elapsed = current_time - last_write_time
|
||
|
||
# 如果两帧间隔太短，考虑合并或跳过
|
||
if elapsed < min_frame_interval and len(recent_frames) >= 2:
|
||
skipped_frames += 1
|
||
continue
|
||
|
||
# 如果两帧间隔太长，进行插值平滑
|
||
if elapsed > max_frame_interval and len(recent_frames) >= 2:
|
||
# 计算需要插入的帧数
|
||
frames_to_insert = min(int(elapsed / target_time_per_frame), 3)
|
||
|
||
for i in range(frames_to_insert):
|
||
# 创建插值帧
|
||
weight = (i + 1) / (frames_to_insert + 1)
|
||
interpolated_frame = cv2.addWeighted(recent_frames[-2], 1 - weight, recent_frames[-1], weight, 0)
|
||
|
||
# 切换双缓冲
|
||
current_buffer = buffer2 if current_buffer is buffer1 else buffer1
|
||
|
||
# 高效调整大小并写入
|
||
interpolated_resized = cv2.resize(interpolated_frame, size, interpolation=cv2.INTER_LINEAR)
|
||
# 在窗口输出图像
|
||
# cv2.imshow("RTMP Stream with YOLO", interpolated_resized)
|
||
|
||
# img_bytes = interpolated_resized.tobytes()
|
||
#
|
||
# # 写入管道
|
||
# pipe.stdin.write(img_bytes)
|
||
# pipe.stdin.flush()
|
||
|
||
# 切换双缓冲
|
||
current_buffer = buffer2 if current_buffer is buffer1 else buffer1
|
||
|
||
# 正常写入当前帧 - 使用高效的调整大小方法
|
||
resized_frame = cv2.resize(frame, size, interpolation=cv2.INTER_LINEAR)
|
||
# cv2.imshow()
|
||
# cv2.imshow("RTMP Stream with YOLO", resized_frame)
|
||
|
||
frame_time = time.time_ns()
|
||
date_str = datetime.datetime.now().strftime("%Y%m%d") # 避免变量覆盖
|
||
|
||
local_path = f"saveImg/output-{frame_time}.jpg"
|
||
frame_count = frame_count + 1
|
||
# # 用作本地测试，保存图像；测试逻辑，随机暂时几张
|
||
if frame_count % 10 == 0:
|
||
success = cv2.imwrite(local_path, resized_frame, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
|
||
if success:
|
||
print(f"图像已保存到: {frame_count}")
|
||
# 上传文件到 Minio
|
||
minio_path, file_type = upload_file(local_path, date_str)
|
||
minio_path, file_type = upload_file(local_path, date_str)
|
||
|
||
message = {
|
||
"task_id": task_id,
|
||
"minio": {
|
||
"minio_path": minio_path,
|
||
"file_type": file_type
|
||
},
|
||
"box_detail": box_detail,
|
||
"uav_location": local_drc_message
|
||
}
|
||
message_json = json.dumps(message, indent=4, ensure_ascii=False)
|
||
|
||
# mqtt_client.put_message(str(message))
|
||
mqtt_client.publish_message(message_json)
|
||
else:
|
||
print("错误: 图像保存失败")
|
||
|
||
if os.path.exists(local_path):
|
||
try:
|
||
os.remove(local_path)
|
||
print(f"文件 {local_path} 已删除")
|
||
except Exception as e:
|
||
print(f"删除失败: {e}")
|
||
# img_bytes = resized_frame.tobytes()
|
||
# pipe.stdin.write(img_bytes)
|
||
# pipe.stdin.flush()
|
||
#
|
||
# frame_count += 1
|
||
if frame_count % 30 == 0:
|
||
current_fps_time = time.time()
|
||
fps = 30 / (current_fps_time - last_fps_time)
|
||
print(f"推流FPS: {fps:.2f}, 跳过的帧: {skipped_frames}, 队列大小: {processed_frame_queue.qsize()}")
|
||
last_fps_time = current_fps_time
|
||
skipped_frames = 0
|
||
|
||
last_write_time = time.time()
|
||
pipe_error_count = 0 # 成功写入后重置错误计数
|
||
|
||
except queue.Empty:
|
||
# 队列为空且有足够的最近帧时，考虑生成插值帧以保持流畅
|
||
if len(recent_frames) >= 2 and time.time() - last_write_time > target_time_per_frame:
|
||
try:
|
||
# 创建插值帧
|
||
interpolated_frame = cv2.addWeighted(recent_frames[-2], 0.5, recent_frames[-1], 0.5, 0)
|
||
|
||
# 切换双缓冲
|
||
current_buffer = buffer2 if current_buffer is buffer1 else buffer1
|
||
|
||
resized_frame = cv2.resize(interpolated_frame, size, interpolation=cv2.INTER_LINEAR)
|
||
# img_bytes = resized_frame.tobytes()
|
||
# pipe.stdin.write(img_bytes)
|
||
# pipe.stdin.flush()
|
||
last_write_time = time.time()
|
||
except Exception:
|
||
pass
|
||
continue
|
||
except Exception as e:
|
||
print(f"写入帧错误: {e}")
|
||
pipe_error_count += 1
|
||
if pipe_error_count >= max_pipe_errors:
|
||
print("FFmpeg管道错误过多，正在终止进程...")
|
||
stop_event.set() # 主动结束所有线程
|
||
break
|
||
|
||
|
||
#
|
||
# def cls2_find(video_path, m1, cls, confidence):
|
||
# try:
|
||
# ov_model = YOLO(m1)
|
||
#
|
||
# # ---------------------------------MinIO 存储路径（用于后续上传）------------------------------------
|
||
# minio_path = "AIResults"
|
||
#
|
||
# # -------------------------------获取当前日期，用于存储图像目录--------------------------------------
|
||
# date_str = datetime.datetime.now().strftime("%Y%m%d")
|
||
# save_dir = f"{date_str}"
|
||
# if not os.path.exists(save_dir):
|
||
# os.makedirs(save_dir)
|
||
#
|
||
# # 打开视频流
|
||
# cap = cv2.VideoCapture(video_path)
|
||
# if not cap.isOpened():
|
||
# print("Error: Could not open video.")
|
||
# return
|
||
#
|
||
# # 获取视频的帧率 (fps)
|
||
# fps = cap.get(cv2.CAP_PROP_FPS)
|
||
#
|
||
# # -----------------------------根据模型设置类别-------------------------------------------------
|
||
# if m1 == "gdaq.pt": # 仅当使用 gdaq.pt 时，保存类别 2 和 4
|
||
# cls2 = [2, 4]
|
||
# elif m1 == "best.pt": # 仅当使用 best.pt 时，保存类别 0
|
||
# cls2 = [0]
|
||
# else: # 其它模型不保存
|
||
# cls2 = []
|
||
#
|
||
# # ------------------------------------------cls2检测--------------------------------------------
|
||
# skip_frames = int(fps * 10) # 设置跳过帧数为 10 秒
|
||
#
|
||
# while cap.isOpened() and not stop_event.is_set():
|
||
# if skip_frames > 0:
|
||
# skip_frames -= 1 # 逐帧减少
|
||
# cap.grab() # 仅抓取帧，不进行解码
|
||
# continue # 跳过处理
|
||
#
|
||
# ret, frame = cap.read()
|
||
# if not ret:
|
||
# break # 无法读取帧时退出
|
||
#
|
||
# # 目标检测
|
||
# results = ov_model(frame, conf=confidence, classes=cls, show=False)
|
||
#
|
||
# for result in results:
|
||
# for box in result.boxes:
|
||
# cls_index = int(box.cls[0]) # 获取类别索引
|
||
#
|
||
# # 如果检测到的类别在 cls2 里，跳过 10 秒
|
||
# if cls_index in cls2:
|
||
# skip_frames = int(fps * 10) # 设置跳过帧数为 10 秒
|
||
# # upload_and_insert_to_db(frame, ov_model, cls_index, save_dir, minio_path)
|
||
# filename = f"{save_dir}/frame_{int(cap.get(cv2.CAP_PROP_POS_FRAMES))}_cls2.jpg"
|
||
# cv2.imwrite(filename, frame)
|
||
# print(f"保存图像: {filename}")
|
||
#
|
||
# # 绘制检测框
|
||
# x1, y1, x2, y2 = map(int, box.xyxy[0])
|
||
# label = f"{result.names[cls_index]} {box.conf[0]:.2f}"
|
||
# cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
|
||
# cv2.putText(frame, label, (x1, y1 - 10),
|
||
# cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
|
||
# except Exception as e:
|
||
# print(f"cls2_find错误: {e}")
|
||
# finally:
|
||
# if 'cap' in locals() and cap is not None:
|
||
# cap.release()
|
||
|
||
|
||
# 后台计算的主要逻辑，识别图像，进而将事件通过mqtt转发、minio存储
|
||
# 1、识别视频流中的事件，进而将事件的前后帧做存储
|
||
# 2、事件基于mqtt 做转发，mqtt可以包含事件，也可以不包含事件
|
||
def startBackAIVideo2(task_id, video_path,push_url, m1, list_func_id, confidence, para, mqtt_ip, mqtt_port, mqtt_topic,
|
||
model_cls, chinese_label):
|
||
global mqtt_client, local_drc_message
|
||
setIfAI(True)
|
||
|
||
cap = None
|
||
read_thread = None
|
||
process_thread = None
|
||
write_thread = None
|
||
read_mqtt_drc_message = None
|
||
ov_model = None
|
||
|
||
local_drc_message = None # 清空数据
|
||
|
||
try:
|
||
# 设置环境变量，提高YOLO性能
|
||
os.environ["OMP_NUM_THREADS"] = "4"
|
||
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
||
|
||
# 导入必要的库
|
||
import torch
|
||
print(f"PyTorch 可用: {torch.__version__}, CUDA可用: {torch.cuda.is_available()}")
|
||
if torch.cuda.is_available():
|
||
print(f"GPU: {torch.cuda.get_device_name(0)}")
|
||
torch.cuda.empty_cache()
|
||
|
||
# 预加载YOLO模型
|
||
print("预加载YOLO模型...")
|
||
max_retries = 3
|
||
retry_count = 0
|
||
model_params = {}
|
||
# 尝试检测YOLO版本以适配不同版本的参数
|
||
try:
|
||
# 先简单尝试加载模型，不带任何参数
|
||
test_model = YOLO(m1)
|
||
# 如果成功，检查可用的参数
|
||
if hasattr(test_model, "task"):
|
||
model_params["task"] = "detect" # 指定任务类型
|
||
|
||
# 检查是否支持half精度
|
||
if torch.cuda.is_available():
|
||
model_params["half"] = True
|
||
|
||
# 检查是否支持verbose参数
|
||
import inspect
|
||
if "verbose" in inspect.signature(YOLO.__init__).parameters:
|
||
model_params["verbose"] = False
|
||
|
||
print(f"检测到支持的YOLO参数: {model_params}")
|
||
except Exception as e:
|
||
print(f"参数检测失败，将使用默认参数: {e}")
|
||
model_params = {}
|
||
|
||
while retry_count < max_retries:
|
||
try:
|
||
|
||
ov_model = YOLO(m1, **model_params)
|
||
dummy_frame = np.zeros((1080, 1920, 3), dtype=np.uint8)
|
||
for _ in range(3):
|
||
# 伪代码，后续需要关联ai_model_list 的id
|
||
cls = [0]
|
||
ov_model(dummy_frame, classes=cls, conf=confidence, show=False)
|
||
print("YOLO模型加载成功并预热完成")
|
||
break
|
||
except Exception as e:
|
||
retry_count += 1
|
||
print(f"YOLO模型加载失败 (尝试 {retry_count}/{max_retries}): {e}")
|
||
if "got an unexpected keyword argument" in str(e) and model_params:
|
||
param_name = str(e).split("'")[-2]
|
||
if param_name in model_params:
|
||
del model_params[param_name]
|
||
time.sleep(2)
|
||
|
||
if ov_model is None:
|
||
raise Exception("无法加载YOLO模型，达到最大重试次数")
|
||
|
||
# 优化模型推理参数
|
||
ov_model.conf = confidence
|
||
|
||
# 打开视频流
|
||
print(f"正在连接视频流: {video_path}")
|
||
cap = cv2.VideoCapture(video_path, cv2.CAP_FFMPEG)
|
||
cap.set(cv2.CAP_PROP_OPEN_TIMEOUT_MSEC, 60000)
|
||
cap.set(cv2.CAP_PROP_READ_TIMEOUT_MSEC, 50000)
|
||
|
||
if not cap.isOpened():
|
||
print("无法打开视频流，尝试强制释放资源...")
|
||
cap.release()
|
||
time.sleep(2)
|
||
raise Exception(f"无法打开视频流: {video_path}")
|
||
|
||
# 设置全局帧队列大小
|
||
global frame_queue, processed_frame_queue
|
||
frame_queue = queue.Queue(maxsize=80)
|
||
processed_frame_queue = queue.Queue(maxsize=40)
|
||
|
||
# 声明mqtt 对象
|
||
mqtt_client = MQTTClient(broker, port, topic)
|
||
|
||
# 与机场维持消息通讯，进而读取最新一条drc消息
|
||
mqtt_client_drc = MQTTClient(mqtt_ip, mqtt_port, mqtt_topic)
|
||
|
||
# 创建并启动线程
|
||
read_thread = Thread(target=read_frames, args=(cap,push_url, frame_queue))
|
||
# process_frames_tricker
|
||
# 读取队列中的frame，并且对frame做二次计算
|
||
process_thread = Thread(target=process_frames_tricker,
|
||
args=(push_url,frame_queue, processed_frame_queue, ov_model, list_func_id, confidence,para,model_cls,chinese_label))
|
||
# process_thread = Thread(target=multi_model_process_frames,
|
||
# args=(
|
||
# frame_queue, processed_frame_queue, ov_model, list_func_id, confidence, para,
|
||
# model_cls,
|
||
# chinese_label))
|
||
# process_thread = Thread(target=process_frames,
|
||
# args=(model_id, frame_queue, processed_frame_queue, ov_model, cls, confidence))
|
||
|
||
write_thread = Thread(target=write_frames,
|
||
args=(processed_frame_queue, None, (1280, 720), mqtt_client, task_id)) # pipe=None
|
||
|
||
read_mqtt_drc_message = Thread(target=read_drc_mqtt, args=(mqtt_client_drc,)) # pipe=None
|
||
|
||
# consumer = Thread(target=consumer_thread, args=(mqtt_client,))
|
||
|
||
read_thread.daemon = True
|
||
process_thread.daemon = True
|
||
write_thread.daemon = True
|
||
read_mqtt_drc_message.daemon = True
|
||
# consumer.daemon = True
|
||
print("开始处理视频流...")
|
||
read_thread.start()
|
||
process_thread.start()
|
||
write_thread.start()
|
||
read_mqtt_drc_message.start()
|
||
# consumer.start()
|
||
|
||
# 等待线程结束
|
||
while getIfAI() and not stop_event.is_set():
|
||
if not (read_thread.is_alive() and process_thread.is_alive()
|
||
and write_thread.is_alive()
|
||
and read_mqtt_drc_message.is_alive()
|
||
# and consumer.is_alive()
|
||
):
|
||
print(f"read_thread.is_alive() {read_thread.is_alive()} "
|
||
f"process_thread.is_alive() {process_thread.is_alive()} "
|
||
f"write_thread.is_alive() {write_thread.is_alive()} ")
|
||
# f"consumer.is_alive() {consumer.is_alive()} ")
|
||
print("检测到某个线程已停止运行，正在终止所有线程...")
|
||
stop_event.set()
|
||
break
|
||
time.sleep(0.1)
|
||
|
||
except Exception as e:
|
||
print(f"发生错误: {e}")
|
||
finally:
|
||
print("正在清理资源...")
|
||
stop_event.set()
|
||
setIfAI(False)
|
||
# 安全关闭MQTT客户端
|
||
if mqtt_client is not None:
|
||
try:
|
||
mqtt_client.close()
|
||
except Exception as e:
|
||
print(f"关闭MQTT客户端时出错: {str(e)}")
|
||
if (read_thread and process_thread
|
||
and write_thread
|
||
):
|
||
timeout = 3
|
||
start_time = time.time()
|
||
while (read_thread.is_alive() or process_thread.is_alive()
|
||
or write_thread.is_alive()
|
||
or read_mqtt_drc_message.is_alive()
|
||
# or consumer.is_alive()
|
||
) and time.time() - start_time < timeout:
|
||
time.sleep(0.1)
|
||
|
||
if cap is not None:
|
||
try:
|
||
cap.release()
|
||
print("视频捕获资源已释放")
|
||
except:
|
||
pass
|
||
|
||
try:
|
||
cv2.destroyAllWindows()
|
||
except:
|
||
pass
|
||
|
||
print("所有资源已清理完毕")
|
||
|
||
|
||
if __name__ == '__main__':
|
||
sn = "1581F6QAD243C00BP71E"
|
||
# video_path = f"rtmp://112.44.103.230:1935/live/{sn}"
|
||
video_path = "rtmp://112.44.103.230:1935/live/123456"
|
||
# FFmpeg 推流地址
|
||
# rtmp = f"rtmp://112.44.103.230:1935/live/{sn}ai"
|
||
mq_ts = time.time_ns() # 时间戳，标识唯一消息
|
||
try:
|
||
startBackAIVideo(mq_ts, video_path, "best.pt", [0, 1, 2, 3, 4], 0.4)
|
||
time.sleep(60)
|
||
except KeyboardInterrupt:
|
||
print("程序被用户中断")
|
||
stopAIVideo()
|
||
except Exception as e:
|
||
print(f"程序异常: {e}")
|
||
# import cv2
|
||
#
|
||
# print(cv2.getBuildInformation())
|