diff --git a/ws/ws.py b/ws/ws.py index 2ed2dcc..1a6051a 100644 --- a/ws/ws.py +++ b/ws/ws.py @@ -1,15 +1,34 @@ import asyncio import datetime import json +import os from contextlib import asynccontextmanager from typing import Dict, Optional +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" # 关键修正:从core目录向上一级找ocr文件夹 +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 = 60 # 客户端超时阈值(秒) +HEARTBEAT_TIMEOUT = 600 # 客户端超时阈值(秒) WS_ENDPOINT = "/ws" # WebSocket端点路径 # -------------------------- 核心数据结构与全局变量 -------------------------- @@ -113,9 +132,31 @@ async def handle_text_msg(client_ip: str, text: str, conn: ClientConnection): async def handle_binary_msg(client_ip: str, data: bytes): """处理二进制消息(保留扩展入口)""" print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}:收到{len(data)}字节二进制数据") - # 将二进制转化为 numpy 数组 - data_ndarray = np.frombuffer(data, dtype=np.uint8) + # 将二进制数据转换为NumPy数组(uint8类型) + nparr = np.frombuffer(data, np.uint8) + # 解码为图像,返回与cv2.imread相同的格式(BGR通道的ndarray) + img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) + #转存到本地images文件夹下 + # 确保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/{client_ip.replace('.', '_')}_{timestamp}.jpg" + # 保存图像到本地 + + 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}") + else: + print("未检测到任何违规内容") + else: + print(f"无法读取测试图像:") # -------------------------- WebSocket核心端点 -------------------------- @ws_router.websocket(WS_ENDPOINT)