yolo模型识别不到
This commit is contained in:
129
ws/ws.py
129
ws/ws.py
@ -17,17 +17,16 @@ from service.device_action_service import add_device_action
|
||||
from schema.device_action_schema import DeviceActionCreate
|
||||
from core.all import detect, load_model
|
||||
|
||||
# -------------------------- 1. AES 加密解密工具（固定密钥）--------------------------
|
||||
# -------------------------- 1. AES 加密工具（仅用于服务器向客户端发送消息）--------------------------
|
||||
AES_SECRET_KEY = b"jr1vA6tfWMHOYi6UXw67UuO6fdak2rMa" # 约定密钥（32字节）
|
||||
AES_BLOCK_SIZE = 16 # AES固定块大小
|
||||
|
||||
|
||||
def aes_encrypt(plaintext: str) -> dict:
|
||||
"""AES-CBC加密：返回{密文, IV, 算法标识}（均Base64编码）"""
|
||||
"""AES-CBC加密：返回{密文, IV, 算法标识}（均Base64编码）- 仅用于服务器发消息"""
|
||||
try:
|
||||
iv = os.urandom(AES_BLOCK_SIZE) # 随机IV（16字节）
|
||||
cipher = AES.new(AES_SECRET_KEY, AES.MODE_CBC, iv)
|
||||
# 明文填充+加密+Base64编码
|
||||
padded_plaintext = pad(plaintext.encode("utf-8"), AES_BLOCK_SIZE)
|
||||
ciphertext = base64.b64encode(cipher.encrypt(padded_plaintext)).decode("utf-8")
|
||||
iv_base64 = base64.b64encode(iv).decode("utf-8")
|
||||
@ -40,20 +39,6 @@ def aes_encrypt(plaintext: str) -> dict:
|
||||
raise Exception(f"AES加密失败: {str(e)}") from e
|
||||
|
||||
|
||||
def aes_decrypt(encrypted_dict: dict) -> str:
|
||||
"""AES-CBC解密：输入加密字典，返回原始文本"""
|
||||
try:
|
||||
# Base64解码密文和IV
|
||||
ciphertext = base64.b64decode(encrypted_dict["ciphertext"])
|
||||
iv = base64.b64decode(encrypted_dict["iv"])
|
||||
# 解密+去除填充
|
||||
cipher = AES.new(AES_SECRET_KEY, AES.MODE_CBC, iv)
|
||||
decrypted = unpad(cipher.decrypt(ciphertext), AES_BLOCK_SIZE).decode("utf-8")
|
||||
return decrypted
|
||||
except Exception as e:
|
||||
raise Exception(f"AES解密失败: {str(e)}") from e
|
||||
|
||||
|
||||
# -------------------------- 2. 配置常量（保持原有）--------------------------
|
||||
HEARTBEAT_INTERVAL = 30 # 心跳检查间隔（秒）
|
||||
HEARTBEAT_TIMEOUT = 600 # 客户端超时阈值（秒）
|
||||
@ -72,7 +57,7 @@ def get_current_time_file_str() -> str:
|
||||
return datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
|
||||
|
||||
|
||||
# -------------------------- 4. 客户端连接封装（新增消息加密）--------------------------
|
||||
# -------------------------- 4. 客户端连接封装（服务器发消息仍加密，接收消息改明文）--------------------------
|
||||
class ClientConnection:
|
||||
def __init__(self, websocket: WebSocket, client_ip: str):
|
||||
self.websocket = websocket
|
||||
@ -96,28 +81,25 @@ class ClientConnection:
|
||||
return self.consumer_task
|
||||
|
||||
async def send_frame_permit(self):
|
||||
"""发送加密的帧许可信号"""
|
||||
"""发送加密的帧许可信号（服务器→客户端：加密）"""
|
||||
try:
|
||||
# 1. 构建原始消息
|
||||
frame_permit_msg = {
|
||||
"type": "frame",
|
||||
"timestamp": get_current_time_str(),
|
||||
"client_ip": self.client_ip
|
||||
}
|
||||
# 2. AES加密消息
|
||||
encrypted_msg = aes_encrypt(json.dumps(frame_permit_msg))
|
||||
# 3. 发送加密消息
|
||||
encrypted_msg = aes_encrypt(json.dumps(frame_permit_msg)) # 保持加密
|
||||
await self.websocket.send_json(encrypted_msg)
|
||||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 已发送加密帧许可")
|
||||
except Exception as e:
|
||||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 帧许可加密/发送失败 - {str(e)}")
|
||||
|
||||
async def consume_frames(self) -> None:
|
||||
"""消费队列中的帧并处理"""
|
||||
"""消费队列中的明文图像帧并处理"""
|
||||
try:
|
||||
while True:
|
||||
frame_data = await self.frame_queue.get()
|
||||
await self.send_frame_permit() # 发送下一帧许可
|
||||
await self.send_frame_permit() # 回复仍加密
|
||||
try:
|
||||
await self.process_frame(frame_data)
|
||||
finally:
|
||||
@ -128,23 +110,22 @@ class ClientConnection:
|
||||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 帧消费错误 - {str(e)}")
|
||||
|
||||
async def process_frame(self, frame_data: bytes) -> None:
|
||||
"""处理单帧图像（含加密危险通知）"""
|
||||
# 二进制转OpenCV图像
|
||||
"""处理明文图像帧（危险通知仍加密发送）"""
|
||||
# 二进制转OpenCV图像（客户端发的是明文二进制，直接解析）
|
||||
nparr = np.frombuffer(frame_data, np.uint8)
|
||||
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
||||
if img is None:
|
||||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 无法解析图像")
|
||||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 无法解析明文图像")
|
||||
return
|
||||
|
||||
try:
|
||||
# 调用检测函数（client_ip + img 双参数）
|
||||
has_violation, data, detector_type = await asyncio.to_thread(
|
||||
detect, self.client_ip, img
|
||||
)
|
||||
print(
|
||||
f"[{get_current_time_str()}] 客户端{self.client_ip}: 检测结果 - 违规: {has_violation}, 类型: {detector_type}")
|
||||
|
||||
# 处理违规逻辑（发送加密危险通知）
|
||||
# 违规通知：服务器→客户端，仍加密
|
||||
if has_violation:
|
||||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 检测到违规 - {data}")
|
||||
# 违规次数+1
|
||||
@ -154,19 +135,17 @@ class ClientConnection:
|
||||
except Exception as e:
|
||||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 违规次数更新失败 - {str(e)}")
|
||||
|
||||
# 1. 构建原始危险通知
|
||||
# 构建危险通知并加密发送
|
||||
danger_msg = {
|
||||
"type": "danger",
|
||||
"timestamp": get_current_time_str(),
|
||||
"client_ip": self.client_ip,
|
||||
"detail": data
|
||||
}
|
||||
# 2. AES加密通知
|
||||
encrypted_danger_msg = aes_encrypt(json.dumps(danger_msg))
|
||||
# 3. 发送加密通知
|
||||
encrypted_danger_msg = aes_encrypt(json.dumps(danger_msg)) # 保持加密
|
||||
await self.websocket.send_json(encrypted_danger_msg)
|
||||
except Exception as e:
|
||||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 图像处理错误 - {str(e)}")
|
||||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}: 明文图像处理错误 - {str(e)}")
|
||||
|
||||
|
||||
# -------------------------- 5. 全局状态与心跳管理（保持原有）--------------------------
|
||||
@ -178,7 +157,6 @@ async def heartbeat_checker():
|
||||
"""全局心跳检查任务"""
|
||||
while True:
|
||||
current_time = get_current_time_str()
|
||||
# 筛选超时客户端
|
||||
timeout_ips = [ip for ip, conn in connected_clients.items() if not conn.is_alive()]
|
||||
|
||||
if timeout_ips:
|
||||
@ -186,11 +164,9 @@ async def heartbeat_checker():
|
||||
for ip in timeout_ips:
|
||||
try:
|
||||
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="心跳超时")
|
||||
# 标记离线
|
||||
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)
|
||||
@ -205,19 +181,16 @@ async def heartbeat_checker():
|
||||
await asyncio.sleep(HEARTBEAT_INTERVAL)
|
||||
|
||||
|
||||
# -------------------------- 6. 消息处理工具（新增消息解密）--------------------------
|
||||
# -------------------------- 6. 客户端明文消息处理（关键修改：删除解密逻辑）--------------------------
|
||||
async def send_heartbeat_ack(conn: ClientConnection):
|
||||
"""发送加密的心跳确认"""
|
||||
"""发送加密的心跳确认（服务器→客户端：加密）"""
|
||||
try:
|
||||
# 1. 构建原始心跳确认
|
||||
heartbeat_ack_msg = {
|
||||
"type": "heart",
|
||||
"timestamp": get_current_time_str(),
|
||||
"client_ip": conn.client_ip
|
||||
}
|
||||
# 2. AES加密
|
||||
encrypted_msg = aes_encrypt(json.dumps(heartbeat_ack_msg))
|
||||
# 3. 发送
|
||||
encrypted_msg = aes_encrypt(json.dumps(heartbeat_ack_msg)) # 保持加密
|
||||
await conn.websocket.send_json(encrypted_msg)
|
||||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 已发送加密心跳确认")
|
||||
return True
|
||||
@ -228,44 +201,22 @@ async def send_heartbeat_ack(conn: ClientConnection):
|
||||
|
||||
|
||||
async def handle_text_msg(conn: ClientConnection, text: str):
|
||||
"""处理加密的文本消息（如心跳）"""
|
||||
"""处理客户端明文文本消息（如心跳）- 关键修改：无需解密，直接解析JSON"""
|
||||
try:
|
||||
# 1. 解析加密字典
|
||||
encrypted_dict = json.loads(text)
|
||||
# 2. AES解密
|
||||
decrypted_text = aes_decrypt(encrypted_dict)
|
||||
# 3. 解析业务消息
|
||||
msg = json.loads(decrypted_text)
|
||||
|
||||
# 客户端发的是明文JSON，直接解析（删除原解密步骤）
|
||||
msg = json.loads(text)
|
||||
if msg.get("type") == "heart":
|
||||
conn.update_heartbeat()
|
||||
await send_heartbeat_ack(conn)
|
||||
await send_heartbeat_ack(conn) # 服务器回复仍加密
|
||||
else:
|
||||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 未知文本类型（{msg.get('type')}）")
|
||||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 未知明文文本类型（{msg.get('type')}）")
|
||||
except json.JSONDecodeError:
|
||||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 无效JSON格式")
|
||||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 无效JSON格式（明文文本）")
|
||||
except Exception as e:
|
||||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 文本消息解密失败 - {str(e)}")
|
||||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 明文文本消息处理失败 - {str(e)}")
|
||||
|
||||
|
||||
async def handle_binary_msg(conn: ClientConnection, data: str):
|
||||
"""处理加密的图像消息（客户端需先转Base64+加密）"""
|
||||
try:
|
||||
# 1. 解密得到Base64编码的图像
|
||||
encrypted_dict = json.loads(data)
|
||||
decrypted_base64 = aes_decrypt(encrypted_dict)
|
||||
# 2. Base64解码为二进制图像
|
||||
frame_data = base64.b64decode(decrypted_base64)
|
||||
# 3. 加入帧队列
|
||||
conn.frame_queue.put_nowait(frame_data)
|
||||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 解密后图像（{len(frame_data)}字节）入队")
|
||||
except asyncio.QueueFull:
|
||||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 帧队列已满，丢弃数据")
|
||||
except Exception as e:
|
||||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}: 图像消息解密失败 - {str(e)}")
|
||||
|
||||
|
||||
# -------------------------- 7. WebSocket路由与生命周期（保持原有结构）--------------------------
|
||||
# -------------------------- 7. WebSocket路由与生命周期（关键修改：处理明文二进制图像）--------------------------
|
||||
ws_router = APIRouter()
|
||||
|
||||
|
||||
@ -276,7 +227,6 @@ async def lifespan(app: FastAPI):
|
||||
heartbeat_task = asyncio.create_task(heartbeat_checker())
|
||||
print(f"[{get_current_time_str()}] 心跳检查任务启动（ID: {id(heartbeat_task)}）")
|
||||
yield
|
||||
# 关闭时清理
|
||||
if heartbeat_task and not heartbeat_task.done():
|
||||
heartbeat_task.cancel()
|
||||
await heartbeat_task
|
||||
@ -285,8 +235,8 @@ async def lifespan(app: FastAPI):
|
||||
|
||||
@ws_router.websocket(WS_ENDPOINT)
|
||||
async def websocket_endpoint(websocket: WebSocket):
|
||||
"""WebSocket连接处理入口"""
|
||||
load_model() # 加载检测模型（仅一次）
|
||||
"""WebSocket连接处理入口 - 关键修改：接收客户端明文二进制图像"""
|
||||
load_model() # 加载检测模型（建议移到全局，避免重复加载）
|
||||
await websocket.accept()
|
||||
client_ip = websocket.client.host if websocket.client else "unknown_ip"
|
||||
current_time = get_current_time_str()
|
||||
@ -306,8 +256,8 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
# 注册新连接
|
||||
new_conn = ClientConnection(websocket, client_ip)
|
||||
connected_clients[client_ip] = new_conn
|
||||
new_conn.start_consumer() # 启动帧消费
|
||||
await new_conn.send_frame_permit() # 发送首次帧许可
|
||||
new_conn.start_consumer()
|
||||
await new_conn.send_frame_permit() # 首次许可仍加密
|
||||
|
||||
# 标记客户端上线
|
||||
try:
|
||||
@ -321,28 +271,33 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
|
||||
print(f"[{current_time}] 客户端{client_ip}: 连接注册成功，在线数: {len(connected_clients)}")
|
||||
|
||||
# 消息循环（接收客户端消息）
|
||||
# 消息循环：接收客户端明文消息（关键修改）
|
||||
while True:
|
||||
data = await websocket.receive()
|
||||
if "text" in data:
|
||||
# 处理加密文本消息（心跳、客户端指令）
|
||||
# 处理客户端明文文本（如心跳：{"type":"heart",...}）
|
||||
await handle_text_msg(new_conn, data["text"])
|
||||
elif "bytes" in data:
|
||||
# 兼容客户端发送二进制：先转Base64再处理
|
||||
base64_data = base64.b64encode(data["bytes"]).decode("utf-8")
|
||||
await handle_binary_msg(new_conn, base64_data)
|
||||
# 处理客户端明文二进制图像（直接入队，无需解密）
|
||||
frame_data = data["bytes"]
|
||||
try:
|
||||
new_conn.frame_queue.put_nowait(frame_data)
|
||||
print(f"[{get_current_time_str()}] 客户端{client_ip}: 明文图像（{len(frame_data)}字节）入队")
|
||||
except asyncio.QueueFull:
|
||||
print(f"[{get_current_time_str()}] 客户端{client_ip}: 帧队列已满，丢弃数据")
|
||||
except Exception as e:
|
||||
print(f"[{get_current_time_str()}] 客户端{client_ip}: 明文图像处理失败 - {str(e)}")
|
||||
|
||||
except WebSocketDisconnect as e:
|
||||
print(f"[{get_current_time_str()}] 客户端{client_ip}: 主动断开连接（代码: {e.code}）")
|
||||
except Exception as e:
|
||||
print(f"[{get_current_time_str()}] 客户端{client_ip}: 连接异常 - {str(e)[:50]}")
|
||||
finally:
|
||||
# 清理资源（断开后处理）
|
||||
# 清理资源
|
||||
if client_ip in connected_clients:
|
||||
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)
|
||||
@ -352,4 +307,4 @@ async def websocket_endpoint(websocket: WebSocket):
|
||||
except Exception as e:
|
||||
print(f"[{get_current_time_str()}] 客户端{client_ip}: 离线更新失败 - {str(e)}")
|
||||
connected_clients.pop(client_ip, None)
|
||||
print(f"[{get_current_time_str()}] 客户端{client_ip}: 资源已清理，在线数: {len(connected_clients)}")
|
||||
print(f"[{get_current_time_str()}] 客户端{client_ip}: 资源已清理，在线数: {len(connected_clients)}")
|
||||
Reference in New Issue
Block a user