387 lines
16 KiB
Python
387 lines
16 KiB
Python
import asyncio
|
||
import datetime
|
||
import json
|
||
import os
|
||
from contextlib import asynccontextmanager
|
||
from typing import Dict, Optional, AsyncGenerator
|
||
from service.device_service import update_online_status_by_ip, increment_alarm_count_by_ip
|
||
from service.device_action_service import add_device_action
|
||
from schema.device_action_schema import DeviceActionCreate
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from fastapi import WebSocket, APIRouter, WebSocketDisconnect, FastAPI
|
||
from queue import Queue
|
||
from threading import Lock
|
||
|
||
from ocr.model_violation_detector import MultiModelViolationDetector
|
||
|
||
# 配置文件路径（建议实际部署时改为相对路径或环境变量）
|
||
YOLO_MODEL_PATH = r"D:\Git\bin\video\ocr\models\best.pt"
|
||
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"
|
||
|
||
# 模型池配置（根据GPU显存调整，每个模型约占1G显存）
|
||
MODEL_POOL_SIZE = 3 # 最大并发客户端数
|
||
|
||
# 配置常量
|
||
HEARTBEAT_INTERVAL = 30 # 心跳检查间隔（秒）
|
||
HEARTBEAT_TIMEOUT = 600 # 客户端超时阈值（秒）
|
||
WS_ENDPOINT = "/ws" # WebSocket端点路径
|
||
FRAME_QUEUE_SIZE = 1 # 帧队列大小限制
|
||
|
||
|
||
# 模型池实现 - 提前初始化固定数量的模型实例
|
||
class ModelPool:
|
||
def __init__(self, pool_size: int = MODEL_POOL_SIZE):
|
||
self.pool = Queue(maxsize=pool_size)
|
||
self.lock = Lock()
|
||
# 提前初始化模型实例（显存会在此阶段预分配）
|
||
for i in range(pool_size):
|
||
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
|
||
)
|
||
self.pool.put(detector)
|
||
print(f"[{get_current_time_str()}] 模型池初始化：第{i + 1}/{pool_size}个模型加载完成")
|
||
|
||
def get_model(self) -> MultiModelViolationDetector:
|
||
"""从池子里获取模型（阻塞直到有可用实例）"""
|
||
with self.lock:
|
||
return self.pool.get()
|
||
|
||
def return_model(self, detector: MultiModelViolationDetector):
|
||
"""将模型归还给池子"""
|
||
with self.lock:
|
||
self.pool.put(detector)
|
||
|
||
|
||
# 初始化模型池（程序启动时加载所有模型，显存会一次性占用 MODEL_POOL_SIZE * 单模型显存）
|
||
model_pool = ModelPool(pool_size=MODEL_POOL_SIZE)
|
||
|
||
|
||
# 工具函数：获取格式化时间字符串
|
||
def get_current_time_str() -> str:
|
||
return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
|
||
|
||
def get_current_time_file_str() -> str:
|
||
return datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
|
||
|
||
|
||
# 客户端连接封装
|
||
class ClientConnection:
|
||
def __init__(self, websocket: WebSocket, client_ip: str):
|
||
self.websocket = websocket
|
||
self.client_ip = client_ip
|
||
self.last_heartbeat = datetime.datetime.now()
|
||
self.frame_queue = asyncio.Queue(maxsize=FRAME_QUEUE_SIZE)
|
||
self.consumer_task: Optional[asyncio.Task] = None
|
||
|
||
# 从模型池获取专属模型（每个客户端独立占用一个模型实例）
|
||
self.detector = model_pool.get_model()
|
||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：已获取模型池中的模型实例（显存独立）")
|
||
|
||
def update_heartbeat(self):
|
||
"""更新心跳时间"""
|
||
self.last_heartbeat = datetime.datetime.now()
|
||
|
||
def is_alive(self) -> bool:
|
||
"""判断客户端是否存活"""
|
||
timeout = (datetime.datetime.now() - self.last_heartbeat).total_seconds()
|
||
return timeout < HEARTBEAT_TIMEOUT
|
||
|
||
def start_consumer(self):
|
||
"""启动帧消费任务"""
|
||
self.consumer_task = asyncio.create_task(self.consume_frames())
|
||
return self.consumer_task
|
||
|
||
def release_model(self):
|
||
"""客户端断开时归还模型到池"""
|
||
model_pool.return_model(self.detector)
|
||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：模型已归还至模型池（显存可复用）")
|
||
|
||
async def send_frame_permit(self):
|
||
"""发送帧发送许可信号"""
|
||
try:
|
||
frame_permit_msg = {
|
||
"type": "frame",
|
||
"timestamp": get_current_time_str(),
|
||
"client_ip": self.client_ip
|
||
}
|
||
await self.websocket.send_json(frame_permit_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:
|
||
# 1. 从队列取出帧
|
||
frame_data = await self.frame_queue.get()
|
||
|
||
# 2. 立即发送下一帧许可
|
||
await self.send_frame_permit()
|
||
|
||
try:
|
||
# 3. 并行处理帧：用线程池执行AI检测（真正并发）
|
||
await self.process_frame(frame_data)
|
||
finally:
|
||
self.frame_queue.task_done()
|
||
|
||
except asyncio.CancelledError:
|
||
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 process_frame(self, frame_data: bytes) -> None:
|
||
"""处理单帧图像数据（使用客户端专属模型）"""
|
||
# 二进制数据转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}：无法解析图像数据")
|
||
return
|
||
|
||
# 确保图像保存目录存在
|
||
os.makedirs('images', exist_ok=True)
|
||
|
||
# 保存图像
|
||
filename = f"images/{self.client_ip.replace('.', '_')}_{get_current_time_file_str()}.jpg"
|
||
try:
|
||
cv2.imwrite(filename, img)
|
||
print(f"[{get_current_time_str()}] 图像已保存至：{filename}")
|
||
|
||
# 关键修改：使用客户端专属模型 + 线程池并行执行AI检测
|
||
has_violation, violation_type, details = await asyncio.to_thread(
|
||
self.detector.detect_violations, # 客户端独立模型
|
||
img # 输入图像
|
||
)
|
||
|
||
if has_violation:
|
||
print(
|
||
f"[{get_current_time_str()}] 客户端{self.client_ip}：检测到违规 - 类型: {violation_type}, 详情: {details}")
|
||
|
||
# 调用违规次数加一方法
|
||
try:
|
||
await asyncio.to_thread(increment_alarm_count_by_ip, self.client_ip)
|
||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：违规次数已+1")
|
||
except Exception as e:
|
||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：违规次数更新失败 - {str(e)}")
|
||
|
||
# 发送危险通知
|
||
danger_msg = {
|
||
"type": "danger",
|
||
"timestamp": get_current_time_str(),
|
||
"client_ip": self.client_ip
|
||
}
|
||
await self.websocket.send_json(danger_msg)
|
||
else:
|
||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：未检测到违规")
|
||
except Exception as e:
|
||
print(f"[{get_current_time_str()}] 客户端{self.client_ip}：图像处理错误 - {str(e)}")
|
||
|
||
|
||
# 全局状态管理
|
||
connected_clients: Dict[str, ClientConnection] = {}
|
||
client_lock = asyncio.Lock() # 保护connected_clients的锁
|
||
heartbeat_task: Optional[asyncio.Task] = None
|
||
|
||
|
||
# 心跳检查
|
||
async def heartbeat_checker():
|
||
while True:
|
||
current_time = get_current_time_str()
|
||
# 加锁保护字典遍历
|
||
async with client_lock:
|
||
timeout_ips = [ip for ip, conn in connected_clients.items() if not conn.is_alive()]
|
||
|
||
if timeout_ips:
|
||
print(f"[{current_time}] 心跳检查：{len(timeout_ips)}个客户端超时（IP：{timeout_ips}）")
|
||
for ip in timeout_ips:
|
||
try:
|
||
async with client_lock:
|
||
conn = connected_clients.get(ip)
|
||
if not conn:
|
||
continue
|
||
|
||
if conn.consumer_task and not conn.consumer_task.done():
|
||
conn.consumer_task.cancel()
|
||
await conn.websocket.close(code=1008, reason="心跳超时")
|
||
|
||
# 归还模型
|
||
conn.release_model()
|
||
|
||
# 超时设为离线并记录
|
||
try:
|
||
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)
|
||
print(f"[{current_time}] 客户端{ip}：已标记为离线并记录操作")
|
||
except Exception as e:
|
||
print(f"[{current_time}] 客户端{ip}：离线状态更新失败 - {str(e)}")
|
||
finally:
|
||
async with client_lock:
|
||
connected_clients.pop(ip, None)
|
||
else:
|
||
async with client_lock:
|
||
print(f"[{current_time}] 心跳检查：{len(connected_clients)}个客户端在线")
|
||
|
||
await asyncio.sleep(HEARTBEAT_INTERVAL)
|
||
|
||
|
||
# 应用生命周期管理
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
global heartbeat_task
|
||
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()
|
||
try:
|
||
await heartbeat_task
|
||
print(f"[{get_current_time_str()}] 全局心跳检查任务已取消")
|
||
except asyncio.CancelledError:
|
||
pass
|
||
|
||
|
||
# 消息处理工具函数
|
||
async def send_heartbeat_ack(conn: ClientConnection):
|
||
try:
|
||
heartbeat_ack_msg = {
|
||
"type": "heart",
|
||
"timestamp": get_current_time_str(),
|
||
"client_ip": conn.client_ip
|
||
}
|
||
await conn.websocket.send_json(heartbeat_ack_msg)
|
||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}：已发送心跳确认")
|
||
return True
|
||
except Exception as e:
|
||
async with client_lock:
|
||
connected_clients.pop(conn.client_ip, None)
|
||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}：心跳确认发送失败 - {str(e)}")
|
||
return False
|
||
|
||
|
||
async def handle_text_msg(conn: ClientConnection, text: str):
|
||
try:
|
||
msg = json.loads(text)
|
||
if msg.get("type") == "heart":
|
||
conn.update_heartbeat()
|
||
await send_heartbeat_ack(conn)
|
||
else:
|
||
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文本消息")
|
||
|
||
|
||
async def handle_binary_msg(conn: ClientConnection, data: bytes):
|
||
try:
|
||
conn.frame_queue.put_nowait(data)
|
||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}：图像数据（{len(data)}字节）已加入队列")
|
||
except asyncio.QueueFull:
|
||
print(f"[{get_current_time_str()}] 客户端{conn.client_ip}：帧队列已满，丢弃当前图像数据")
|
||
|
||
|
||
# WebSocket路由配置
|
||
ws_router = APIRouter()
|
||
|
||
|
||
@ws_router.websocket(WS_ENDPOINT)
|
||
async def websocket_endpoint(websocket: WebSocket):
|
||
await websocket.accept()
|
||
client_ip = websocket.client.host if websocket.client else "unknown_ip"
|
||
current_time = get_current_time_str()
|
||
print(f"[{current_time}] 客户端{client_ip}：WebSocket连接已建立")
|
||
|
||
is_online_updated = False
|
||
new_conn = None
|
||
|
||
try:
|
||
# 处理重复连接
|
||
async with client_lock:
|
||
if client_ip in connected_clients:
|
||
old_conn = connected_clients[client_ip]
|
||
if old_conn.consumer_task and not old_conn.consumer_task.done():
|
||
old_conn.consumer_task.cancel()
|
||
await old_conn.websocket.close(code=1008, reason="同一IP新连接建立")
|
||
old_conn.release_model() # 归还旧连接的模型
|
||
connected_clients.pop(client_ip)
|
||
print(f"[{current_time}] 客户端{client_ip}：已关闭旧连接并回收模型")
|
||
|
||
# 注册新连接
|
||
new_conn = ClientConnection(websocket, client_ip)
|
||
async with client_lock:
|
||
connected_clients[client_ip] = new_conn
|
||
new_conn.start_consumer()
|
||
# 初始许可：连接建立后立即发一次
|
||
await new_conn.send_frame_permit()
|
||
|
||
# 标记上线并记录
|
||
try:
|
||
await asyncio.to_thread(update_online_status_by_ip, client_ip, 1)
|
||
action_data = DeviceActionCreate(client_ip=client_ip, action=1)
|
||
await asyncio.to_thread(add_device_action, action_data)
|
||
print(f"[{current_time}] 客户端{client_ip}：已标记为在线并记录操作")
|
||
is_online_updated = True
|
||
except Exception as e:
|
||
print(f"[{current_time}] 客户端{client_ip}：上线状态更新失败 - {str(e)}")
|
||
|
||
async with client_lock:
|
||
print(f"[{current_time}] 客户端{client_ip}：新连接注册成功，在线数：{len(connected_clients)}")
|
||
|
||
# 消息循环
|
||
while True:
|
||
data = await websocket.receive()
|
||
if "text" in data:
|
||
await handle_text_msg(new_conn, data["text"])
|
||
elif "bytes" in data:
|
||
await handle_binary_msg(new_conn, data["bytes"])
|
||
|
||
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 new_conn and client_ip in connected_clients:
|
||
async with client_lock:
|
||
conn = connected_clients.get(client_ip)
|
||
if conn:
|
||
if conn.consumer_task and not conn.consumer_task.done():
|
||
conn.consumer_task.cancel()
|
||
|
||
# 归还模型到模型池
|
||
conn.release_model()
|
||
|
||
# 主动/异常断开时标记离线
|
||
if is_online_updated:
|
||
try:
|
||
await asyncio.to_thread(update_online_status_by_ip, client_ip, 0)
|
||
action_data = DeviceActionCreate(client_ip=client_ip, action=0)
|
||
await asyncio.to_thread(add_device_action, action_data)
|
||
print(f"[{get_current_time_str()}] 客户端{client_ip}：断开后已标记为离线")
|
||
except Exception as e:
|
||
print(f"[{get_current_time_str()}] 客户端{client_ip}：断开后离线更新失败 - {str(e)}")
|
||
|
||
connected_clients.pop(client_ip, None)
|
||
|
||
async with client_lock:
|
||
print(f"[{get_current_time_str()}] 客户端{client_ip}：资源已清理，在线数：{len(connected_clients)}")
|
||
|
||
|
||
# 创建FastAPI应用
|
||
app = FastAPI(lifespan=lifespan)
|
||
app.include_router(ws_router)
|
||
|
||
if __name__ == "__main__":
|
||
import uvicorn
|
||
|
||
uvicorn.run(app, host="0.0.0.0", port=8000)
|