diff --git a/ws/ws.py b/ws/ws.py index 6be7b33..8efca08 100644 --- a/ws/ws.py +++ b/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)}个") \ No newline at end of file