zzx/video
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

优化

Browse Source
This commit is contained in:
ZZX9599
2025-09-04 12:29:27 +08:00
parent ea82a33a8f
commit b5d870a19c
5 changed files with 461 additions and 339 deletions
Download Patch File Download Diff File Expand all files Collapse all files

309
ws/ws.py
View File

@ -4,6 +4,9 @@ import json
import os
from contextlib import asynccontextmanager
from typing import Dict, Optional, AsyncGenerator
from service.device_service import update_online_status_by_ip, increment_alarm_count_by_ip
from service.device_action_service import add_device_action
from schema.device_action_schema import DeviceActionCreate
import cv2
import numpy as np
@ -11,7 +14,7 @@ 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"
@ -26,265 +29,287 @@ detector = MultiModelViolationDetector(
ocr_confidence_threshold=0.5
)
# -------------------------- 配置常量 --------------------------
# 配置常量
HEARTBEAT_INTERVAL = 30 # 心跳检查间隔（秒）
HEARTBEAT_TIMEOUT = 600 # 客户端超时阈值（秒）
WS_ENDPOINT = "/ws" # WebSocket端点路径
FRAME_QUEUE_SIZE = 1 # 帧队列大小限制（保持1，确保单帧处理）
# -------------------------- 核心数据结构与全局变量 --------------------------
ws_router = APIRouter()
FRAME_QUEUE_SIZE = 1 # 帧队列大小限制
# 客户端连接封装（包含帧队列）
# 工具函数：获取格式化时间字符串（统一时间戳格式）
def get_current_time_str() -> str:
return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
def get_current_time_file_str() -> str:
return datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
# 客户端连接封装
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 # 消费者任务
self.frame_queue = asyncio.Queue(maxsize=FRAME_QUEUE_SIZE)
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 send_allow_send_frame(self):
"""向客户端发送JSON信号，通知其可发送下一帧二进制数据"""
async def send_frame_permit(self):
"""
发送「帧发送许可信号」
通知客户端可发送下一帧图像
"""
try:
allow_msg = {
"type": "allow_send_frame", # 信号类型，与客户端约定
"timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
"status": "ready", # 表示服务器已准备好接收下一帧
"client_ip": self.client_ip # 可选：便于客户端确认自身身份
frame_permit_msg = {
"type": "frame",
"timestamp": get_current_time_str(),
"client_ip": self.client_ip
}
await self.websocket.send_json(allow_msg)
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：已发送「允许发送帧」信号")
await self.websocket.send_json(frame_permit_msg)
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：已发送帧发送许可信号（取帧后立即通知）")
except Exception as e:
# 发送失败大概率是客户端已断开，不影响主流程，仅日志记录
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：发送「允许发送帧」信号失败 - {str(e)}")
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：帧许可信号发送失败 - {str(e)}")
# 帧消费协程
async def consume_frames(self) -> None:
"""从队列中获取帧并进行处理，处理完后通知客户端可发送下一帧"""
"""消费队列中的帧并处理（核心调整：取帧后立即发许可，再处理帧）"""
try:
while True:
# 从队列获取帧数据（队列空时会阻塞，等待客户端发送）
# 1. 从队列取出帧（阻塞直到有帧可用）
frame_data = await self.frame_queue.get()
# -------------------------- 核心修改：取出帧后立即发送下一帧许可 --------------------------
await self.send_frame_permit() # 取帧即通知客户端发下一帧，无需等处理完成
# -----------------------------------------------------------------------------------------
try:
# 处理帧数据
# 2. 处理取出的帧（即使处理慢，客户端也已收到许可，可提前准备下一帧）
await self.process_frame(frame_data)
finally:
# 标记任务完成（队列计数-1，此时队列回到空状态）
# 3. 标记帧任务完成（无论处理成功/失败，都需清理队列）
self.frame_queue.task_done()
# ---------- 修改：处理完当前帧后，立即通知客户端可发送下一帧 ----------
await self.send_allow_send_frame()
except asyncio.CancelledError:
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：帧消费任务已取消")
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：帧消费任务已取消")
except Exception as e:
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：帧处理错误 - {str(e)}")
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：帧消费逻辑错误 - {str(e)}")
async def process_frame(self, frame_data: bytes) -> None:
"""处理单帧图像数据（原有逻辑不变）"""
# 将二进制数据转换为NumPy数组（uint8类型）
"""处理单帧图像数据（检测违规后发送危险通知 + 调用违规次数加一方法）"""
# 二进制数据转OpenCV图像
nparr = np.frombuffer(frame_data, np.uint8)
# 解码为图像，返回与cv2.imread相同的格式（BGR通道的ndarray）
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
if img is None:
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：无法解析图像数据")
return
# 确保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"
# 确保图像保存目录存在
os.makedirs('images', exist_ok=True)
# 保存图像（按IP+时间戳命名，避免冲突）
filename = f"images/{self.client_ip.replace('.', '_')}_{get_current_time_file_str()}.jpg"
try:
# 保存图像到本地
cv2.imwrite(filename, img)
print(f"[{datetime.datetime.now():%H:%M:%S}] 图像已保存至：{filename}")
print(f"[{get_current_time_str()}] 图像已保存至：{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(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：未检测到任何违规内容")
# 执行违规检测
has_violation, violation_type, details = detector.detect_violations(img)
if has_violation:
print(
f"[{get_current_time_str()}] 客户端{self.client_ip}：检测到违规 - 类型: {violation_type}, 详情: {details}")
# 调用违规次数加一方法
try:
await asyncio.to_thread(increment_alarm_count_by_ip, self.client_ip)
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：违规次数已+1")
except Exception as e:
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：违规次数更新失败 - {str(e)}")
# 发送「危险通知」
danger_msg = {
"type": "danger",
"timestamp": get_current_time_str(),
"client_ip": self.client_ip
}
await self.websocket.send_json(danger_msg)
else:
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：无法解析图像数据")
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：未检测到违规")
except Exception as e:
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：图像处理错误 - {str(e)}")
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：图像处理错误 - {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. 筛选超时客户端（避免遍历中修改字典）
current_time = get_current_time_str()
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}）")
print(f"[{current_time}] 心跳检查：{len(timeout_ips)}个客户端超时（IP：{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="心跳超时")
conn = connected_clients[ip]
if conn.consumer_task and not conn.consumer_task.done():
conn.consumer_task.cancel()
await conn.websocket.close(code=1008, reason="心跳超时")
# 超时设为离线并记录
try:
await asyncio.to_thread(update_online_status_by_ip, ip, 0)
action_data = DeviceActionCreate(client_ip=ip, action=0)
await asyncio.to_thread(add_device_action, action_data)
print(f"[{current_time}] 客户端{ip}：已标记为离线并记录操作")
except Exception as e:
print(f"[{current_time}] 客户端{ip}：离线状态更新失败 - {str(e)}")
finally:
connected_clients.pop(ip, None)
else:
print(f"[{now:%H:%M:%S}] 心跳检查：{len(connected_clients)}个客户端在线，无超时")
print(f"[{current_time}] 心跳检查：{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)}）")
print(f"[{get_current_time_str()}] 全局心跳检查任务启动（任务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}] 心跳任务已取消")
print(f"[{get_current_time_str()}] 全局心跳检查任务已取消")
except asyncio.CancelledError:
pass
# -------------------------- 消息处理（文本/心跳逻辑不变，二进制逻辑保留） --------------------------
async def send_heartbeat_ack(client_ip: str):
"""回复心跳确认（原有逻辑不变）"""
if client_ip not in connected_clients:
return False
# 消息处理工具函数
async def send_heartbeat_ack(conn: ClientConnection):
try:
ack = {
"timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
"type": "heartbeat"
heartbeat_ack_msg = {
"type": "heart",
"timestamp": get_current_time_str(),
"client_ip": conn.client_ip
}
await connected_clients[client_ip].websocket.send_json(ack)
await conn.websocket.send_json(heartbeat_ack_msg)
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}：已发送心跳确认")
return True
except Exception:
connected_clients.pop(client_ip, None)
except Exception as e:
connected_clients.pop(conn.client_ip, None)
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}：心跳确认发送失败 - {str(e)}")
return False
async def handle_text_msg(client_ip: str, text: str, conn: ClientConnection):
"""处理文本消息（核心：心跳+JSON解析，原有逻辑不变）"""
async def handle_text_msg(conn: ClientConnection, text: str):
try:
msg = json.loads(text)
# 仅处理心跳类型消息
if msg.get("type") == "heartbeat":
if msg.get("type") == "heart":
conn.update_heartbeat()
await send_heartbeat_ack(client_ip)
await send_heartbeat_ack(conn)
else:
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：收到文本消息：{msg}")
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}：未知文本消息类型（{msg.get('type')}）")
except json.JSONDecodeError:
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：无效JSON消息")
print(f"[{get_current_time_str()}] 客户端{conn.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
# 队列未满，添加帧到队列
async def handle_binary_msg(conn: ClientConnection, data: bytes):
try:
conn.frame_queue.put_nowait(data)
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：已接收{len(data)}字节二进制数据，加入队列")
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}：图像数据（{len(data)}字节）已加入队列")
except asyncio.QueueFull:
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：队列突然满了，丢弃{len(data)}字节数据")
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}：帧队列已满，丢弃当前图像数据")
# WebSocket路由配置
ws_router = APIRouter()
# -------------------------- 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}：连接成功")
client_ip = websocket.client.host if websocket.client else "unknown_ip"
current_time = get_current_time_str()
print(f"[{current_time}] 客户端{client_ip}：WebSocket连接已建立")
is_online_updated = False
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新连接")
old_conn = connected_clients[client_ip]
if old_conn.consumer_task and not old_conn.consumer_task.done():
old_conn.consumer_task.cancel()
await old_conn.websocket.close(code=1008, reason="同一IP新连接建立")
connected_clients.pop(client_ip)
print(f"[{now:%H:%M:%S}] 客户端{client_ip}：关闭旧连接")
print(f"[{current_time}] 客户端{client_ip}：已关闭旧连接")
# 注册新连接
new_conn = ClientConnection(websocket, client_ip)
connected_clients[client_ip] = new_conn
new_conn.start_consumer()
# 初始许可：连接建立后立即发一次，让客户端知道可发第一帧（后续靠取帧后自动发）
await new_conn.send_frame_permit()
# 启动帧消费任务
consumer_task = new_conn.start_consumer()
# ---------- 修改：客户端刚连接时，队列空，立即发送「允许发送帧」信号 ----------
await new_conn.send_allow_send_frame()
print(f"[{now:%H:%M:%S}] 客户端{client_ip}：注册成功，已启动帧消费任务，当前在线{len(connected_clients)}个")
# 标记上线并记录
try:
await asyncio.to_thread(update_online_status_by_ip, client_ip, 1)
action_data = DeviceActionCreate(client_ip=client_ip, action=1)
await asyncio.to_thread(add_device_action, action_data)
print(f"[{current_time}] 客户端{client_ip}：已标记为在线并记录操作")
is_online_updated = True
except Exception as e:
print(f"[{current_time}] 客户端{client_ip}：上线状态更新失败 - {str(e)}")
# 循环接收消息（原有逻辑不变）
print(f"[{current_time}] 客户端{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)
await handle_text_msg(new_conn, data["text"])
elif "bytes" in data:
await handle_binary_msg(client_ip, data["bytes"])
await handle_binary_msg(new_conn, data["bytes"])
# 异常处理（断开/错误）
except WebSocketDisconnect as e:
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：主动断开（代码：{e.code}）")
print(f"[{get_current_time_str()}] 客户端{client_ip}：主动断开连接（代码：{e.code}）")
except Exception as e:
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：连接异常（{str(e)[:50]}）")
print(f"[{get_current_time_str()}] 客户端{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()
conn = connected_clients[client_ip]
if conn.consumer_task and not conn.consumer_task.done():
conn.consumer_task.cancel()
# 主动/异常断开时标记离线
if is_online_updated:
try:
await asyncio.to_thread(update_online_status_by_ip, client_ip, 0)
action_data = DeviceActionCreate(client_ip=client_ip, action=0)
await asyncio.to_thread(add_device_action, action_data)
print(f"[{get_current_time_str()}] 客户端{client_ip}：断开后已标记为离线")
except Exception as e:
print(f"[{get_current_time_str()}] 客户端{client_ip}：断开后离线更新失败 - {str(e)}")
connected_clients.pop(client_ip, None)
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：连接已清理，当前在线{len(connected_clients)}个")
print(f"[{get_current_time_str()}] 客户端{client_ip}：资源已清理，在线数：{len(connected_clients)}")