WS方案
This commit is contained in:
ZZX9599
65
ws/ws.py
65
ws/ws.py
@ -30,7 +30,7 @@ detector = MultiModelViolationDetector(
|
||||
HEARTBEAT_INTERVAL = 30 # 心跳检查间隔（秒）
|
||||
HEARTBEAT_TIMEOUT = 600 # 客户端超时阈值（秒）
|
||||
WS_ENDPOINT = "/ws" # WebSocket端点路径
|
||||
FRAME_QUEUE_SIZE = 1 # 帧队列大小限制
|
||||
FRAME_QUEUE_SIZE = 1 # 帧队列大小限制（保持1，确保单帧处理）
|
||||
|
||||
# -------------------------- 核心数据结构与全局变量 --------------------------
|
||||
ws_router = APIRouter()
|
||||
@ -59,26 +59,44 @@ class ClientConnection:
|
||||
self.consumer_task = asyncio.create_task(self.consume_frames())
|
||||
return self.consumer_task
|
||||
|
||||
# ---------- 新增：发送“允许发送二进制帧”的信号给客户端 ----------
|
||||
async def send_allow_send_frame(self):
|
||||
"""向客户端发送JSON信号，通知其可发送下一帧二进制数据"""
|
||||
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 # 可选：便于客户端确认自身身份
|
||||
}
|
||||
await self.websocket.send_json(allow_msg)
|
||||
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 consume_frames(self) -> None:
|
||||
"""从队列中获取帧并进行处理"""
|
||||
"""从队列中获取帧并进行处理，处理完后通知客户端可发送下一帧"""
|
||||
try:
|
||||
while True:
|
||||
# 从队列获取帧数据
|
||||
# 从队列获取帧数据（队列空时会阻塞，等待客户端发送）
|
||||
frame_data = await self.frame_queue.get()
|
||||
try:
|
||||
# 处理帧数据
|
||||
await self.process_frame(frame_data)
|
||||
finally:
|
||||
# 标记任务完成
|
||||
# 标记任务完成（队列计数-1，此时队列回到空状态）
|
||||
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}：帧消费任务已取消")
|
||||
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）
|
||||
@ -102,7 +120,7 @@ class ClientConnection:
|
||||
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,
|
||||
@ -111,11 +129,11 @@ class ClientConnection:
|
||||
"timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
})
|
||||
else:
|
||||
print("未检测到任何违规内容")
|
||||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：未检测到任何违规内容")
|
||||
else:
|
||||
print(f"无法解析图像数据")
|
||||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：无法解析图像数据")
|
||||
except Exception as e:
|
||||
print(f"[{datetime.datetime.now():%H:%M:%S}] 图像处理错误 - {str(e)}")
|
||||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：图像处理错误 - {str(e)}")
|
||||
|
||||
|
||||
# 全局连接管理（IP -> 连接实例）
|
||||
@ -124,7 +142,7 @@ connected_clients: Dict[str, ClientConnection] = {}
|
||||
heartbeat_task: Optional[asyncio.Task] = None
|
||||
|
||||
|
||||
# -------------------------- 心跳检查逻辑 --------------------------
|
||||
# -------------------------- 心跳检查逻辑（原有逻辑不变） --------------------------
|
||||
async def heartbeat_checker():
|
||||
while True:
|
||||
now = datetime.datetime.now()
|
||||
@ -149,7 +167,7 @@ async def heartbeat_checker():
|
||||
await asyncio.sleep(HEARTBEAT_INTERVAL)
|
||||
|
||||
|
||||
# -------------------------- 应用生命周期 --------------------------
|
||||
# -------------------------- 应用生命周期（原有逻辑不变） --------------------------
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
global heartbeat_task
|
||||
@ -167,9 +185,9 @@ async def lifespan(app: FastAPI):
|
||||
pass
|
||||
|
||||
|
||||
# -------------------------- 消息处理 --------------------------
|
||||
# -------------------------- 消息处理（文本/心跳逻辑不变，二进制逻辑保留） --------------------------
|
||||
async def send_heartbeat_ack(client_ip: str):
|
||||
"""回复心跳确认"""
|
||||
"""回复心跳确认（原有逻辑不变）"""
|
||||
if client_ip not in connected_clients:
|
||||
return False
|
||||
try:
|
||||
@ -185,7 +203,7 @@ async def send_heartbeat_ack(client_ip: str):
|
||||
|
||||
|
||||
async def handle_text_msg(client_ip: str, text: str, conn: ClientConnection):
|
||||
"""处理文本消息（核心：心跳+JSON解析）"""
|
||||
"""处理文本消息（核心：心跳+JSON解析，原有逻辑不变）"""
|
||||
try:
|
||||
msg = json.loads(text)
|
||||
# 仅处理心跳类型消息
|
||||
@ -193,36 +211,33 @@ async def handle_text_msg(client_ip: str, text: str, conn: ClientConnection):
|
||||
conn.update_heartbeat()
|
||||
await send_heartbeat_ack(client_ip)
|
||||
else:
|
||||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：收到消息：{msg}")
|
||||
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)}字节数据到队列")
|
||||
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核心端点 --------------------------
|
||||
# -------------------------- WebSocket核心端点（修改连接初始化逻辑） --------------------------
|
||||
@ws_router.websocket(WS_ENDPOINT)
|
||||
async def websocket_endpoint(websocket: WebSocket):
|
||||
# 接受连接 + 获取客户端IP
|
||||
@ -248,9 +263,11 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
|
||||
# 启动帧消费任务
|
||||
consumer_task = new_conn.start_consumer()
|
||||
# ---------- 修改：客户端刚连接时，队列空，立即发送「允许发送帧」信号 ----------
|
||||
await new_conn.send_allow_send_frame()
|
||||
print(f"[{now:%H:%M:%S}] 客户端{client_ip}：注册成功，已启动帧消费任务，当前在线{len(connected_clients)}个")
|
||||
|
||||
# 循环接收消息
|
||||
# 循环接收消息（原有逻辑不变）
|
||||
while True:
|
||||
data = await websocket.receive()
|
||||
if "text" in data:
|
||||
@ -270,4 +287,4 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
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)}个")
|
||||
print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：连接已清理，当前在线{len(connected_clients)}个")
|
||||
Reference in New Issue
Block a user