274 lines
11 KiB
Python
274 lines
11 KiB
Python
import asyncio
|
||
import datetime
|
||
import json
|
||
import os
|
||
from contextlib import asynccontextmanager
|
||
from typing import Dict, Optional, AsyncGenerator
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from fastapi import WebSocket, APIRouter, WebSocketDisconnect, FastAPI
|
||
|
||
from ocr.model_violation_detector import MultiModelViolationDetector
|
||
|
||
# 配置文件相对路径（根据实际目录结构调整）
|
||
YOLO_MODEL_PATH = r"D:\Git\bin\video\ocr\models\best.pt"
|
||
FORBIDDEN_WORDS_PATH = r"D:\Git\bin\video\ocr\forbidden_words.txt"
|
||
OCR_CONFIG_PATH = r"D:\Git\bin\video\ocr\config\1.yaml"
|
||
KNOWN_FACES_DIR = r"D:\Git\bin\video\ocr\known_faces"
|
||
|
||
# 创建检测器实例
|
||
detector = MultiModelViolationDetector(
|
||
forbidden_words_path=FORBIDDEN_WORDS_PATH,
|
||
ocr_config_path=OCR_CONFIG_PATH,
|
||
yolo_model_path=YOLO_MODEL_PATH,
|
||
known_faces_dir=KNOWN_FACES_DIR,
|
||
ocr_confidence_threshold=0.5
|
||
)
|
||
|
||
# -------------------------- 配置常量 --------------------------
|
||
HEARTBEAT_INTERVAL = 30 # 心跳检查间隔（秒）
|
||
HEARTBEAT_TIMEOUT = 600 # 客户端超时阈值（秒）
|
||
WS_ENDPOINT = "/ws" # WebSocket端点路径
|
||
FRAME_QUEUE_SIZE = 1 # 帧队列大小限制
|
||
|
||
# -------------------------- 核心数据结构与全局变量 --------------------------
|
||
ws_router = APIRouter()
|
||
|
||
|
||
# 客户端连接封装（包含帧队列）
|
||
class ClientConnection:
|
||
def __init__(self, websocket: WebSocket, client_ip: str):
|
||
self.websocket = websocket
|
||
self.client_ip = client_ip
|
||
self.last_heartbeat = datetime.datetime.now()
|
||
self.frame_queue = asyncio.Queue(maxsize=FRAME_QUEUE_SIZE) # 帧队列，长度为1
|
||
self.consumer_task: Optional[asyncio.Task] = None # 消费者任务
|
||
|
||
# 更新心跳时间
|
||
def update_heartbeat(self):
|
||
self.last_heartbeat = datetime.datetime.now()
|
||
|
||
# 检查是否存活（超时返回False）
|
||
def is_alive(self) -> bool:
|
||
timeout = (datetime.datetime.now() - self.last_heartbeat).total_seconds()
|
||
return timeout < HEARTBEAT_TIMEOUT
|
||
|
||
# 启动帧消费任务
|
||
def start_consumer(self):
|
||
self.consumer_task = asyncio.create_task(self.consume_frames())
|
||
return self.consumer_task
|
||
|
||
# 帧消费协程
|
||
async def consume_frames(self) -> None:
|
||
"""从队列中获取帧并进行处理"""
|
||
try:
|
||
while True:
|
||
# 从队列获取帧数据
|
||
frame_data = await self.frame_queue.get()
|
||
try:
|
||
# 处理帧数据
|
||
await self.process_frame(frame_data)
|
||
finally:
|
||
# 标记任务完成
|
||
self.frame_queue.task_done()
|
||
except asyncio.CancelledError:
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：帧消费任务已取消")
|
||
except Exception as e:
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：帧处理错误 - {str(e)}")
|
||
|
||
async def process_frame(self, frame_data: bytes) -> None:
|
||
"""处理单帧图像数据"""
|
||
# 将二进制数据转换为NumPy数组（uint8类型）
|
||
nparr = np.frombuffer(frame_data, np.uint8)
|
||
# 解码为图像，返回与cv2.imread相同的格式（BGR通道的ndarray）
|
||
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
||
|
||
# 确保images文件夹存在
|
||
if not os.path.exists('images'):
|
||
os.makedirs('images')
|
||
|
||
# 生成唯一的文件名，包含时间戳和客户端IP，避免文件名冲突
|
||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
|
||
filename = f"images/{self.client_ip.replace('.', '_')}_{timestamp}.jpg"
|
||
|
||
try:
|
||
# 保存图像到本地
|
||
cv2.imwrite(filename, img)
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 图像已保存至：{filename}")
|
||
|
||
# 进行检测
|
||
if img is not None:
|
||
has_violation, violation_type, details = detector.detect_violations(img)
|
||
if has_violation:
|
||
print(f"检测到违规 - 类型: {violation_type}, 详情: {details}")
|
||
# 可以在这里添加发送检测结果回客户端的逻辑
|
||
await self.websocket.send_json({
|
||
"type": "detection_result",
|
||
"has_violation": has_violation,
|
||
"violation_type": violation_type,
|
||
"details": details,
|
||
"timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
})
|
||
else:
|
||
print("未检测到任何违规内容")
|
||
else:
|
||
print(f"无法解析图像数据")
|
||
except Exception as e:
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 图像处理错误 - {str(e)}")
|
||
|
||
|
||
# 全局连接管理（IP -> 连接实例）
|
||
connected_clients: Dict[str, ClientConnection] = {}
|
||
# 心跳任务（全局引用，用于关闭时清理）
|
||
heartbeat_task: Optional[asyncio.Task] = None
|
||
|
||
|
||
# -------------------------- 心跳检查逻辑 --------------------------
|
||
async def heartbeat_checker():
|
||
while True:
|
||
now = datetime.datetime.now()
|
||
# 1. 筛选超时客户端（避免遍历中修改字典）
|
||
timeout_ips = [ip for ip, conn in connected_clients.items() if not conn.is_alive()]
|
||
|
||
# 2. 处理超时连接（关闭+移除）
|
||
if timeout_ips:
|
||
print(f"[{now:%H:%M:%S}] 心跳检查：{len(timeout_ips)}个客户端超时（{timeout_ips}）")
|
||
for ip in timeout_ips:
|
||
try:
|
||
# 取消消费者任务
|
||
if connected_clients[ip].consumer_task and not connected_clients[ip].consumer_task.done():
|
||
connected_clients[ip].consumer_task.cancel()
|
||
await connected_clients[ip].websocket.close(code=1008, reason="心跳超时")
|
||
finally:
|
||
connected_clients.pop(ip, None)
|
||
else:
|
||
print(f"[{now:%H:%M:%S}] 心跳检查：{len(connected_clients)}个客户端在线，无超时")
|
||
|
||
# 3. 等待下一轮检查
|
||
await asyncio.sleep(HEARTBEAT_INTERVAL)
|
||
|
||
|
||
# -------------------------- 应用生命周期 --------------------------
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
global heartbeat_task
|
||
# 启动心跳任务
|
||
heartbeat_task = asyncio.create_task(heartbeat_checker())
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 心跳任务启动（ID：{id(heartbeat_task)}）")
|
||
yield
|
||
# 关闭时取消心跳任务
|
||
if heartbeat_task and not heartbeat_task.done():
|
||
heartbeat_task.cancel()
|
||
try:
|
||
await heartbeat_task
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 心跳任务已取消")
|
||
except asyncio.CancelledError:
|
||
pass
|
||
|
||
|
||
# -------------------------- 消息处理 --------------------------
|
||
async def send_heartbeat_ack(client_ip: str):
|
||
"""回复心跳确认"""
|
||
if client_ip not in connected_clients:
|
||
return False
|
||
try:
|
||
ack = {
|
||
"timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||
"type": "heartbeat"
|
||
}
|
||
await connected_clients[client_ip].websocket.send_json(ack)
|
||
return True
|
||
except Exception:
|
||
connected_clients.pop(client_ip, None)
|
||
return False
|
||
|
||
|
||
async def handle_text_msg(client_ip: str, text: str, conn: ClientConnection):
|
||
"""处理文本消息（核心：心跳+JSON解析）"""
|
||
try:
|
||
msg = json.loads(text)
|
||
# 仅处理心跳类型消息
|
||
if msg.get("type") == "heartbeat":
|
||
conn.update_heartbeat()
|
||
await send_heartbeat_ack(client_ip)
|
||
else:
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：收到消息：{msg}")
|
||
except json.JSONDecodeError:
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：无效JSON消息")
|
||
|
||
|
||
async def handle_binary_msg(client_ip: str, data: bytes):
|
||
"""处理二进制消息（使用队列控制帧处理）"""
|
||
if client_ip not in connected_clients:
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：连接不存在，丢弃{len(data)}字节数据")
|
||
return
|
||
|
||
conn = connected_clients[client_ip]
|
||
|
||
# 检查队列是否已满
|
||
if conn.frame_queue.full():
|
||
# 队列已满，丢弃当前帧
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：队列已满，丢弃{len(data)}字节数据")
|
||
return
|
||
|
||
# 队列未满，添加帧到队列
|
||
try:
|
||
# 非阻塞添加（因为已检查队列未满，所以不会阻塞）
|
||
conn.frame_queue.put_nowait(data)
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：已添加{len(data)}字节数据到队列")
|
||
except asyncio.QueueFull:
|
||
# 理论上不会走到这里，因为上面已检查
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：队列突然满了，丢弃{len(data)}字节数据")
|
||
|
||
|
||
# -------------------------- WebSocket核心端点 --------------------------
|
||
@ws_router.websocket(WS_ENDPOINT)
|
||
async def websocket_endpoint(websocket: WebSocket):
|
||
# 接受连接 + 获取客户端IP
|
||
await websocket.accept()
|
||
client_ip = websocket.client.host if websocket.client else "unknown"
|
||
now = datetime.datetime.now()
|
||
print(f"[{now:%H:%M:%S}] 客户端{client_ip}：连接成功")
|
||
|
||
consumer_task = None
|
||
try:
|
||
# 处理重复连接（关闭旧连接）
|
||
if client_ip in connected_clients:
|
||
# 取消旧连接的消费者任务
|
||
if connected_clients[client_ip].consumer_task and not connected_clients[client_ip].consumer_task.done():
|
||
connected_clients[client_ip].consumer_task.cancel()
|
||
await connected_clients[client_ip].websocket.close(code=1008, reason="同一IP新连接")
|
||
connected_clients.pop(client_ip)
|
||
print(f"[{now:%H:%M:%S}] 客户端{client_ip}：关闭旧连接")
|
||
|
||
# 注册新连接
|
||
new_conn = ClientConnection(websocket, client_ip)
|
||
connected_clients[client_ip] = new_conn
|
||
|
||
# 启动帧消费任务
|
||
consumer_task = new_conn.start_consumer()
|
||
print(f"[{now:%H:%M:%S}] 客户端{client_ip}：注册成功，已启动帧消费任务，当前在线{len(connected_clients)}个")
|
||
|
||
# 循环接收消息
|
||
while True:
|
||
data = await websocket.receive()
|
||
if "text" in data:
|
||
await handle_text_msg(client_ip, data["text"], new_conn)
|
||
elif "bytes" in data:
|
||
await handle_binary_msg(client_ip, data["bytes"])
|
||
|
||
# 异常处理（断开/错误）
|
||
except WebSocketDisconnect as e:
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：主动断开（代码：{e.code}）")
|
||
except Exception as e:
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：连接异常（{str(e)[:50]}）")
|
||
finally:
|
||
# 清理连接和任务
|
||
if client_ip in connected_clients:
|
||
# 取消消费者任务
|
||
if connected_clients[client_ip].consumer_task and not connected_clients[client_ip].consumer_task.done():
|
||
connected_clients[client_ip].consumer_task.cancel()
|
||
connected_clients.pop(client_ip, None)
|
||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：连接已清理，当前在线{len(connected_clients)}个")
|