import base64 import requests as http_client import queue import time # 确保 time 被导入，如果之前被误删 import cv2 import traceback import threading # 确保导入 threading import av # 重新导入 av import numpy as np # from fastapi import requests # 从 rfdetr_core 导入 RFDETRDetector 仅用于类型提示 (可选) from rfdetr_core import RFDETRDetector # 目标检测处理函数 # 函数签名已更改：现在接受一个 detector_instance 作为参数 def yolo_frame(rtc_q: queue.Queue, yolo_q: queue.Queue, stream_detector_instance: RFDETRDetector): thread_name = threading.current_thread().name # 获取线程名称用于日志 print(f"处理线程 '{thread_name}' 已启动。") error_message_displayed_once = False no_detector_message_displayed_once = False # 用于只提示一次没有检测器 if stream_detector_instance is None: print(f"错误 (线程 '{thread_name}'): 未提供有效的检测器实例给yolo_frame。线程将无法处理视频。") # 此线程实际上无法做任何有用的工作，可以考虑直接退出或进入一个安全循环 # 为简单起见，我们允许它进入主循环，但它会在每次迭代时打印警告 while True: frame = None # current_category_counts = {} # 将在获取后转换 try: # 恢复队列长度打印 print(f"线程 '{thread_name}' - 原始队列长度: {rtc_q.qsize()}, 检测队列长度: {yolo_q.qsize()}") frame = rtc_q.get(timeout=0.1) if frame is None: print(f"处理线程 '{thread_name}' 接收到停止信号，正在退出...") # 发送包含None frame和空计数的字典作为停止信号 yolo_q.put({"frame": None, "category_counts": {}}) break category_counts_for_packet = {} if stream_detector_instance: no_detector_message_displayed_once = False # 检测器有效，重置提示 annotated_frame = stream_detector_instance.detect_and_draw_count(frame) error_message_displayed_once = False # 获取英文键的类别计数 english_counts = stream_detector_instance.category_counts.copy() if hasattr(stream_detector_instance, 'category_counts') else {} # 转换为中文键的类别计数 if hasattr(stream_detector_instance, 'VISDRONE_CLASSES_CHINESE'): chinese_map = stream_detector_instance.VISDRONE_CLASSES_CHINESE for eng_key, count_val in english_counts.items(): # 使用 get 提供一个默认值，以防某个英文类别在中文映射表中确实没有 chi_key = chinese_map.get(eng_key, eng_key) category_counts_for_packet[chi_key] = count_val else: # 如果没有中文映射表，则直接使用英文计数 (或记录警告) category_counts_for_packet = english_counts # logger.warning(f"线程 '{thread_name}': stream_detector_instance 没有 VISDRONE_CLASSES_CHINESE 属性，将使用英文类别计数。") else: # 如果没有有效的检测器实例传递进来 if not no_detector_message_displayed_once: print(f"警告 (线程 '{thread_name}'): 无有效检测器实例。将在帧上绘制提示。") no_detector_message_displayed_once = True annotated_frame = frame.copy() cv2.putText(annotated_frame, "No detector instance provided for this stream", (30, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA) category_counts_for_packet = {} # 无检测器，计数为空 # 将帧和类别计数一起放入队列 data_packet = {"frame": annotated_frame, "category_counts": category_counts_for_packet} try: yolo_q.put_nowait(data_packet) except queue.Full: # print(f"警告 (线程 '{thread_name}'): yolo_q 已满，丢弃帧。") # 避免刷屏 pass except queue.Empty: time.sleep(0.01) continue except Exception as e: if not error_message_displayed_once: print(f"线程 '{thread_name}' (yolo_frame) 处理时发生严重错误: {e}") traceback.print_exc() error_message_displayed_once = True time.sleep(1) if frame is not None: try: pass except queue.Full: pass continue print(f"处理线程 '{thread_name}' 已停止。") def push_frame(yolo_q: queue.Queue, rtmp_url: str, gateway: str, frequency: int, push_url: str): thread_name = threading.current_thread().name print(f"推流线程 '{thread_name}' (RTMP: {rtmp_url}) 已启动。") output_container = None stream = None first_frame_processed = False last_push_time = 0 # 记录上次推送base64的时间 try: while True: frame_to_push = None received_category_counts = {} # 初始化为空字典 try: data_packet = yolo_q.get(timeout=0.1) if data_packet: frame_to_push = data_packet.get("frame") received_category_counts = data_packet.get("category_counts", {}) else: # data_packet is None (不太可能，除非队列明确放入None) time.sleep(0.01) continue except queue.Empty: time.sleep(0.01) continue if frame_to_push is None: # 这是通过 data_packet["frame"] is None 来判断的 print(f"推流线程 '{thread_name}' 接收到停止信号，正在清理并退出...") break if not first_frame_processed: if frame_to_push is not None: try: height, width, _ = frame_to_push.shape print(f"线程 '{thread_name}': 首帧尺寸 {width}x{height}，正在初始化RTMP推流器到 {rtmp_url}") output_container = av.open(rtmp_url, 'w', format='flv') stream = output_container.add_stream('libx264', rate=25) stream.pix_fmt = 'yuv420p' stream.width = width stream.height = height stream.options = {'preset': 'ultrafast', 'tune': 'zerolatency', 'crf': '25'} print(f"线程 '{thread_name}': RTMP推流器初始化成功。") first_frame_processed = True except Exception as e_init: print(f"错误 (线程 '{thread_name}'): 初始化PyAV推流容器/流失败: {e_init}") traceback.print_exc() return else: continue if not output_container or not stream: print(f"错误 (线程 '{thread_name}'): 推流器未初始化，无法推流。可能是首帧处理失败。") time.sleep(1) continue # 持续推流到RTMP try: video_frame = av.VideoFrame.from_ndarray(frame_to_push, format='bgr24') for packet in stream.encode(video_frame): output_container.mux(packet) except Exception as e_push: print(f"错误 (线程 '{thread_name}'): 推送帧到RTMP时发生错误: {e_push}") time.sleep(0.5) # 定时推送base64帧 current_time = time.time() if current_time - last_push_time >= frequency: # 将接收到的类别计数传递给 push_base64_frame push_base64_frame(frame_to_push, gateway, push_url, thread_name, received_category_counts) last_push_time = current_time except Exception as e_outer: print(f"推流线程 '{thread_name}' 发生严重外部错误: {e_outer}") traceback.print_exc() finally: print(f"推流线程 '{thread_name}': 进入finally块，准备关闭推流器。") if stream and output_container: try: print(f"推流线程 '{thread_name}': 正在编码流的剩余部分...") for packet in stream.encode(None): output_container.mux(packet) print(f"推流线程 '{thread_name}': 编码剩余部分完成。") except Exception as e_flush: print(f"错误 (线程 '{thread_name}'): 关闭推流流时发生编码/刷新错误: {e_flush}") traceback.print_exc() if output_container: try: print(f"推流线程 '{thread_name}': 正在关闭推流容器...") output_container.close() print(f"推流线程 '{thread_name}': 推流容器已关闭。") except Exception as e_close: print(f"错误 (线程 '{thread_name}'): 关闭推流容器时发生错误: {e_close}") traceback.print_exc() print(f"推流线程 '{thread_name}' 已停止并完成清理。") def push_base64_frame(frame: np.ndarray, gateway: str, push_url: str, thread_name: str, category_counts: dict): """将帧转换为base64并推送到指定URL""" try: # 转换为JPEG格式 _, buffer = cv2.imencode('.jpg', frame) # 转换为base64字符串 frame_base64 = base64.b64encode(buffer).decode('utf-8') # 构建JSON数据 data = { "gateway": gateway, "frame_base64": frame_base64, "category_counts": category_counts # 使用传入的 category_counts } print(f"DEBUG push_base64_frame: Pushing data: {data.get('category_counts')}") # 调试打印，检查发送的数据 # 发送POST请求 response = http_client.post(push_url, json=data, timeout=5) # 检查响应 if response.status_code == 200: print(f"线程 '{thread_name}': base64帧已成功推送到 {push_url}") else: print(f"错误 (线程 '{thread_name}'): 推送base64帧失败，状态码: {response.status_code}") except Exception as e: print(f"错误 (线程 '{thread_name}'): 处理或推送base64帧时发生错误: {e}")