WS方案

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)}个")