173 lines
6.4 KiB
Python
173 lines
6.4 KiB
Python
import asyncio
|
||
from datetime import datetime
|
||
|
||
import aiohttp
|
||
import cv2
|
||
import numpy as np
|
||
from aiortc import RTCPeerConnection, RTCSessionDescription, RTCConfiguration
|
||
from ocr.ocr_violation_detector import OCRViolationDetector
|
||
from ws.ws import send_message_to_client
|
||
|
||
|
||
async def rtc_frame_receiver(url, frame_queue, stop_event):
|
||
"""
|
||
接收RTC帧并往队列放入cv2格式的帧数据
|
||
当队列已满时直接丢弃新帧，不阻塞等待
|
||
当stop_event被设置时停止接收
|
||
"""
|
||
pc = RTCPeerConnection(RTCConfiguration(iceServers=[]))
|
||
|
||
# 累计帧计数器和丢弃帧计数器
|
||
total_frames = 0
|
||
dropped_frames = 0
|
||
|
||
@pc.on("track")
|
||
async def on_track(track):
|
||
nonlocal total_frames, dropped_frames
|
||
if track.kind == "video":
|
||
print("接收到视频轨道、开始接收视频帧")
|
||
while not stop_event.is_set(): # 检查是否需要停止
|
||
# 接收当前帧并累计计数
|
||
frame = await track.recv()
|
||
total_frames += 1
|
||
|
||
# 转换为cv2兼容的BGR格式numpy数组
|
||
frame_cv2 = frame.to_ndarray(format='bgr24')
|
||
|
||
# 验证是否为cv2兼容格式
|
||
if isinstance(frame_cv2, np.ndarray) and frame_cv2.ndim == 3 and frame_cv2.shape[2] == 3:
|
||
# 检查队列是否已满
|
||
if frame_queue.full():
|
||
# 队列已满，丢弃当前帧
|
||
dropped_frames += 1
|
||
print(f"第{total_frames}帧：队列已满，丢弃该帧（累计丢弃: {dropped_frames}）")
|
||
else:
|
||
# 队列未满，放入当前帧
|
||
await frame_queue.put(frame_cv2)
|
||
print(f"第{total_frames}帧：已放入队列")
|
||
else:
|
||
print("帧格式转换失败，不是有效的cv2格式")
|
||
|
||
# 创建并设置本地offer
|
||
offer = await pc.createOffer()
|
||
print("已创建本地 SDP Offer")
|
||
await pc.setLocalDescription(offer)
|
||
|
||
# 发送offer到服务器
|
||
async with aiohttp.ClientSession() as session:
|
||
print("开始向服务器发送 SDP Offer")
|
||
try:
|
||
async with session.post(
|
||
url,
|
||
data=offer.sdp.encode(),
|
||
headers={
|
||
"Content-Type": "application/sdp",
|
||
"Content-Length": str(len(offer.sdp))
|
||
},
|
||
ssl=False
|
||
) as response:
|
||
if response.status == 200:
|
||
print("已接收到服务器的响应、开始处理 SDP Answer")
|
||
answer_sdp = await response.text()
|
||
await pc.setRemoteDescription(RTCSessionDescription(sdp=answer_sdp, type='answer'))
|
||
else:
|
||
print(f"服务器响应错误: {response.status}")
|
||
stop_event.set()
|
||
except Exception as e:
|
||
print(f"发送SDP Offer失败: {str(e)}")
|
||
stop_event.set()
|
||
|
||
try:
|
||
# 保持连接，直到收到停止信号
|
||
while not stop_event.is_set():
|
||
await asyncio.sleep(1)
|
||
except KeyboardInterrupt:
|
||
print("用户中断")
|
||
finally:
|
||
print(f"开始关闭 RTCPeerConnection，共接收{total_frames}帧，丢弃{dropped_frames}帧")
|
||
await pc.close()
|
||
print("已关闭 RTCPeerConnection")
|
||
|
||
|
||
async def frame_consumer(ip, frame_queue, stop_event):
|
||
"""
|
||
从队列中阻塞读取cv2帧并处理（队列为空时阻塞等待）
|
||
检测到违规内容后设置stop_event以终止所有任务
|
||
"""
|
||
# 创建OCR检测器实例
|
||
ocr_detector = OCRViolationDetector(
|
||
forbidden_words_path=r"D:\Git\bin\video\ocr\forbidden_words.txt",
|
||
ocr_confidence_threshold=0.5, )
|
||
|
||
while not stop_event.is_set(): # 检查是否需要停止
|
||
try:
|
||
# 阻塞等待队列中的帧
|
||
current_frame = await frame_queue.get()
|
||
|
||
# 进行OCR检测
|
||
has_violation, words, confidences = ocr_detector.detect(current_frame)
|
||
print(f"检测结果: {'有违规内容' if has_violation else '无违规内容'}")
|
||
print(f"检测到的词: {words}")
|
||
print(f"置信度: {confidences}")
|
||
|
||
# 输出所有检测到的违禁词
|
||
if has_violation:
|
||
print(f"测试结果：图片中共检测到 {len(words)} 个违禁词：")
|
||
response_data = {
|
||
"status": "stop",
|
||
"timestamp": datetime.now().isoformat(),
|
||
"violations": [{"word": w, "confidence": c} for w, c in zip(words, confidences)]
|
||
}
|
||
await send_message_to_client(ip, response_data)
|
||
|
||
for word, conf in zip(words, confidences):
|
||
print(f"- {word}（置信度：{conf:.4f}）")
|
||
|
||
# 检测到违规，设置停止事件
|
||
print("检测到违规内容，准备关闭AI检测")
|
||
stop_event.set()
|
||
|
||
# 标记任务完成，允许生产者放入新的帧
|
||
frame_queue.task_done()
|
||
|
||
except Exception as e:
|
||
print(f"处理帧时发生错误: {str(e)}")
|
||
frame_queue.task_done()
|
||
|
||
|
||
def process_webrtc_stream(ip, webrtc_url):
|
||
"""
|
||
处理WEBRTC流并持续打印OCR检测结果，检测到违规后关闭
|
||
队列大小为1，满时直接丢弃新帧
|
||
|
||
Args:
|
||
ip: IP地址
|
||
webrtc_url: WEBRTC服务器地址
|
||
"""
|
||
# 创建队列（大小为1）和停止事件
|
||
frame_queue = asyncio.Queue(maxsize=1) # 只存储一帧
|
||
stop_event = asyncio.Event() # 用于控制任务停止的事件
|
||
|
||
# 定义事件循环中的主任务
|
||
async def main_task():
|
||
# 创建任务
|
||
receiver_task = asyncio.create_task(rtc_frame_receiver(webrtc_url, frame_queue, stop_event))
|
||
consumer_task = asyncio.create_task(frame_consumer(ip, frame_queue, stop_event))
|
||
|
||
# 等待任务完成
|
||
await asyncio.gather(receiver_task, consumer_task)
|
||
|
||
# 确保队列处理完毕
|
||
await frame_queue.join()
|
||
|
||
try:
|
||
# 运行事件循环
|
||
asyncio.run(main_task())
|
||
except KeyboardInterrupt:
|
||
print("用户中断处理流程")
|
||
stop_event.set()
|
||
finally:
|
||
# 确保关闭所有cv2窗口
|
||
cv2.destroyAllWindows()
|
||
print("AI检测已关闭")
|