diff --git a/schema/device_action_schema.py b/schema/device_action_schema.py new file mode 100644 index 0000000..783d09e --- /dev/null +++ b/schema/device_action_schema.py @@ -0,0 +1,36 @@ +from datetime import datetime +from typing import Optional, List +from pydantic import BaseModel, Field + + +# ------------------------------ +# 请求模型（新增记录用，极简） +# ------------------------------ +class DeviceActionCreate(BaseModel): + """设备操作记录创建模型（0=离线，1=上线）""" + client_ip: str = Field(..., description="客户端IP") + action: int = Field(..., ge=0, le=1, description="操作状态（0=离线，1=上线）") + + +# ------------------------------ +# 响应模型（单条记录） +# ------------------------------ +class DeviceActionResponse(BaseModel): + """设备操作记录响应模型（与自增表对齐）""" + id: int = Field(..., description="自增主键ID") + client_ip: Optional[str] = Field(None, description="客户端IP") + action: Optional[int] = Field(None, description="操作状态（0=离线，1=上线）") + created_at: datetime = Field(..., description="记录创建时间") + updated_at: datetime = Field(..., description="记录更新时间") + + # 支持从数据库结果直接转换 + model_config = {"from_attributes": True} + + +# ------------------------------ +# 列表响应模型（仅含 total + device_actions） +# ------------------------------ +class DeviceActionListResponse(BaseModel): + """设备操作记录列表（仅核心返回字段）""" + total: int = Field(..., description="总记录数") + device_actions: List[DeviceActionResponse] = Field(..., description="操作记录列表") \ No newline at end of file diff --git a/schema/device_schema.py b/schema/device_schema.py index bf80632..adf7e53 100644 --- a/schema/device_schema.py +++ b/schema/device_schema.py @@ -1,4 +1,3 @@ -import hashlib from datetime import datetime from typing import Optional, List, Dict @@ -6,42 +5,31 @@ from pydantic import BaseModel, Field # ------------------------------ -# 请求模型（前端传参校验） +# 请求模型 # ------------------------------ class DeviceCreateRequest(BaseModel): - """设备流信息创建请求模型""" + """设备流信息创建请求模型（与数据库表字段对齐）""" ip: Optional[str] = Field(..., max_length=100, description="设备IP地址") hostname: Optional[str] = Field(None, max_length=100, description="设备别名") - params: Optional[Dict] = Field(None, description="设备详细信息") - - -def md5_encrypt(text: str) -> str: - """对字符串进行MD5加密""" - if not text: - return "" - md5_hash = hashlib.md5() - md5_hash.update(text.encode('utf-8')) - return md5_hash.hexdigest() + params: Optional[Dict] = Field(None, description="设备详细信息（JSON格式）") # ------------------------------ -# 响应模型（后端返回设备数据） +# 响应模型（后端返回数据）- 严格对齐数据库表字段 # ------------------------------ class DeviceResponse(BaseModel): - """设备流信息响应模型（字段与表结构完全对齐）""" - id: int = Field(..., description="设备ID") + """设备流信息响应模型（与数据库表字段完全一致）""" + id: int = Field(..., description="设备主键ID") + client_ip: Optional[str] = Field(None, max_length=100, description="设备IP地址") hostname: Optional[str] = Field(None, max_length=100, description="设备别名") - rtmp_push_url: Optional[str] = Field(None, description="需要推送的RTMP地址") - live_webrtc_url: Optional[str] = Field(None, description="直播的Webrtc地址") - detection_webrtc_url: Optional[str] = Field(None, description="检测的Webrtc地址") device_online_status: int = Field(..., description="设备在线状态（1-在线、0-离线）") device_type: Optional[str] = Field(None, description="设备类型") alarm_count: int = Field(..., description="报警次数") - params: Optional[str] = Field(None, description="设备详细信息") + params: Optional[str] = Field(None, description="设备详细信息（JSON字符串）") created_at: datetime = Field(..., description="记录创建时间") updated_at: datetime = Field(..., description="记录更新时间") - # 支持从数据库查询结果转换 + # 支持从数据库查询结果直接转换 model_config = {"from_attributes": True} diff --git a/service/device_action_service.py b/service/device_action_service.py new file mode 100644 index 0000000..29acb73 --- /dev/null +++ b/service/device_action_service.py @@ -0,0 +1,119 @@ +from fastapi import APIRouter, Query +from mysql.connector import Error as MySQLError + +from ds.db import db +from schema.device_action_schema import ( + DeviceActionCreate, + DeviceActionResponse, + DeviceActionListResponse +) +from schema.response_schema import APIResponse + + +# 路由配置 +router = APIRouter( + prefix="/device/actions", + tags=["设备操作记录"] +) + + +# ------------------------------ +# 内部方法：新增设备操作记录（适配id自增） +# ------------------------------ +def add_device_action(action_data: DeviceActionCreate) -> DeviceActionResponse: + """ + 新增设备操作记录（内部方法，非接口） + :param action_data: 含client_ip和action（0/1） + :return: 新增的完整记录 + """ + conn = None + cursor = None + try: + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + + # 插入SQL（id自增，依赖数据库自动生成） + insert_query = """ + INSERT INTO device_action + (client_ip, action, created_at, updated_at) + VALUES (%s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + """ + cursor.execute(insert_query, ( + action_data.client_ip, + action_data.action + )) + conn.commit() + + # 获取新增记录（通过自增ID查询） + new_id = cursor.lastrowid + cursor.execute("SELECT * FROM device_action WHERE id = %s", (new_id,)) + new_action = cursor.fetchone() + + return DeviceActionResponse(**new_action) + + except MySQLError as e: + if conn: + conn.rollback() + raise Exception(f"新增记录失败：{str(e)}") from e + finally: + db.close_connection(conn, cursor) + + +# ------------------------------ +# 接口：分页查询操作记录列表（仅返回 total + device_actions） +# ------------------------------ +@router.get("/list", response_model=APIResponse, summary="分页查询设备操作记录") +async def get_device_action_list( + page: int = Query(1, ge=1, description="页码，默认1"), + page_size: int = Query(10, ge=1, le=100, description="每页条数，1-100"), + client_ip: str = Query(None, description="按客户端IP筛选"), + action: int = Query(None, ge=0, le=1, description="按状态筛选（0=离线，1=上线）") +): + conn = None + cursor = None + try: + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + + # 1. 构建筛选条件（参数化查询，避免注入） + where_clause = [] + params = [] + if client_ip: + where_clause.append("client_ip = %s") + params.append(client_ip) + if action is not None: + where_clause.append("action = %s") + params.append(action) + + # 2. 查询总记录数（用于返回 total） + count_sql = "SELECT COUNT(*) AS total FROM device_action" + if where_clause: + count_sql += " WHERE " + " AND ".join(where_clause) + cursor.execute(count_sql, params) + total = cursor.fetchone()["total"] + + # 3. 分页查询记录（按创建时间倒序，确保最新记录在前） + offset = (page - 1) * page_size + list_sql = "SELECT * FROM device_action" + if where_clause: + list_sql += " WHERE " + " AND ".join(where_clause) + list_sql += " ORDER BY created_at DESC LIMIT %s OFFSET %s" + params.extend([page_size, offset]) # 追加分页参数（page/page_size仅用于查询，不返回） + + cursor.execute(list_sql, params) + action_list = cursor.fetchall() + + # 4. 仅返回 total + device_actions + return APIResponse( + code=200, + message="查询成功", + data=DeviceActionListResponse( + total=total, + device_actions=[DeviceActionResponse(**item) for item in action_list] + ) + ) + + except MySQLError as e: + raise Exception(f"查询记录失败：{str(e)}") from e + finally: + db.close_connection(conn, cursor) \ No newline at end of file diff --git a/service/device_service.py b/service/device_service.py index 798f4b7..dde502d 100644 --- a/service/device_service.py +++ b/service/device_service.py @@ -1,25 +1,11 @@ import json -import threading -import time -from fastapi import HTTPException, Query, APIRouter, Depends, Request +from fastapi import APIRouter, Query, HTTPException from mysql.connector import Error as MySQLError -from ds.config import LIVE_CONFIG from ds.db import db -from middle.auth_middleware import get_current_user -# 注意：导入的Schema已更新字段 -from schema.device_schema import ( - DeviceCreateRequest, - DeviceResponse, - DeviceListResponse, - md5_encrypt -) +from schema.device_schema import DeviceCreateRequest, DeviceResponse, DeviceListResponse from schema.response_schema import APIResponse -from schema.user_schema import UserResponse - -# 导入之前封装的WEBRTC处理函数 -from core.rtmp import rtmp_pull_video_stream router = APIRouter( prefix="/devices", @@ -27,65 +13,128 @@ router = APIRouter( ) -# 在后台线程中运行WEBRTC处理 -def run_webrtc_processing(ip, webrtc_url): - try: - print(f"开始处理来自设备 {ip} 的WEBRTC流: {webrtc_url}") - rtmp_pull_video_stream(webrtc_url) - except Exception as e: - print(f"WEBRTC处理出错: {str(e)}") - - # ------------------------------ -# 1. 创建设备信息 +# 内部工具方法 - 通过客户端IP增加设备报警次数 # ------------------------------ -@router.post("/add", response_model=APIResponse, summary="创建设备信息") -async def create_device(request: Request, device_data: DeviceCreateRequest): +def increment_alarm_count_by_ip(client_ip: str) -> bool: + """ + 通过客户端IP增加设备的报警次数（内部服务方法） + + :param client_ip: 客户端IP地址 + :return: 操作是否成功 + """ + if not client_ip: + raise ValueError("客户端IP不能为空") + + conn = None + cursor = None + try: + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + + # 检查设备是否存在 + cursor.execute("SELECT id FROM devices WHERE client_ip = %s", (client_ip,)) + if not cursor.fetchone(): + raise ValueError(f"客户端IP为 {client_ip} 的设备不存在") + + # 报警次数加1，并更新时间戳 + update_query = """ + UPDATE devices + SET alarm_count = alarm_count + 1, + updated_at = CURRENT_TIMESTAMP + WHERE client_ip = %s + """ + cursor.execute(update_query, (client_ip,)) + conn.commit() + + return True + except MySQLError as e: + if conn: + conn.rollback() + raise Exception(f"更新报警次数失败：{str(e)}") from e + finally: + db.close_connection(conn, cursor) + + +# ------------------------------ +# 内部工具方法 - 通过客户端IP更新设备在线状态 +# ------------------------------ +def update_online_status_by_ip(client_ip: str, online_status: int) -> bool: + """ + 通过客户端IP更新设备的在线状态（内部服务方法） + + :param client_ip: 客户端IP地址 + :param online_status: 在线状态（1-在线、0-离线） + :return: 操作是否成功 + """ + if not client_ip: + raise ValueError("客户端IP不能为空") + + # 验证状态值有效性 + if online_status not in (0, 1): + raise ValueError("在线状态必须是0（离线）或1（在线）") + + conn = None + cursor = None + try: + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + + # 检查设备是否存在 + cursor.execute("SELECT id FROM devices WHERE client_ip = %s", (client_ip,)) + if not cursor.fetchone(): + raise ValueError(f"客户端IP为 {client_ip} 的设备不存在") + + # 更新在线状态和时间戳 + update_query = """ + UPDATE devices + SET device_online_status = %s, + updated_at = CURRENT_TIMESTAMP + WHERE client_ip = %s + """ + cursor.execute(update_query, (online_status, client_ip)) + conn.commit() + + return True + except MySQLError as e: + if conn: + conn.rollback() + raise Exception(f"更新设备在线状态失败：{str(e)}") from e + finally: + db.close_connection(conn, cursor) + + +# ------------------------------ +# 原有接口保持不变 +# ------------------------------ +@router.post("/add", response_model=APIResponse, summary="创建设备信息") +async def create_device(device_data: DeviceCreateRequest): + # 原有代码保持不变 conn = None cursor = None try: conn = db.get_connection() cursor = conn.cursor(dictionary=True) - # 检查client_ip是否已存在 cursor.execute("SELECT * FROM devices WHERE client_ip = %s", (device_data.ip,)) existing_device = cursor.fetchone() if existing_device: - # 设备创建成功后，在后台线程启动WEBRTC流处理 - threading.Thread( - target=run_webrtc_processing, - # args=(device_data.ip, existing_device["live_webrtc_url"]), - args=(device_data.ip, existing_device["rtmp_push_url"]), - - daemon=True # 设为守护线程，主程序退出时自动结束 - ).start() - # IP已存在时返回该设备信息 + # 更新设备状态为在线 + update_online_status_by_ip(client_ip=device_data.ip, online_status=1) + # 返回信息 return APIResponse( code=200, - message=f"客户端IP {device_data.ip} 已存在", + message=f"设备IP {device_data.ip} 已存在，返回已有设备信息", data=DeviceResponse(**existing_device) ) - # 获取RTMP URL和WEBRTC URL配置 - rtmp_url = str(LIVE_CONFIG.get("rtmp_url", "")) - webrtc_url = str(LIVE_CONFIG.get("webrtc_url", "")) - - # 将设备详细信息（params）转换为JSON字符串 - device_params_json = json.dumps(device_data.params) if device_data.params else None - - # 对JSON字符串进行MD5加密 - device_md5 = md5_encrypt(device_params_json) if device_params_json else "" - - # 解析User-Agent获取设备类型 + from fastapi import Request + request = Request(scope={"type": "http"}) user_agent = request.headers.get("User-Agent", "").lower() - # 优先处理User-Agent为default的情况 if user_agent == "default": - # 检查params中是否存在os键 - if device_data.params and isinstance(device_data.params, dict) and "os" in device_data.params: - device_type = device_data.params["os"] - else: - device_type = "unknown" + device_type = device_data.params.get("os") if ( + device_data.params and isinstance(device_data.params, dict)) else "unknown" elif "windows" in user_agent: device_type = "windows" elif "android" in user_agent: @@ -95,22 +144,16 @@ async def create_device(request: Request, device_data: DeviceCreateRequest): else: device_type = "unknown" - # 构建完整的WEBRTC URL - full_webrtc_url = webrtc_url + device_md5 + device_params_json = json.dumps(device_data.params) if device_data.params else None - # SQL插入语句 insert_query = """ INSERT INTO devices - (client_ip, hostname, rtmp_push_url, live_webrtc_url, detection_webrtc_url, - device_online_status, device_type, alarm_count, params) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) + (client_ip, hostname, device_online_status, device_type, alarm_count, params) + VALUES (%s, %s, %s, %s, %s, %s) """ cursor.execute(insert_query, ( device_data.ip, device_data.hostname, - rtmp_url + device_md5, - full_webrtc_url, # 存储完整的WEBRTC URL - "", 1, device_type, 0, @@ -118,28 +161,22 @@ async def create_device(request: Request, device_data: DeviceCreateRequest): )) conn.commit() - # 获取刚创建的设备信息 device_id = cursor.lastrowid cursor.execute("SELECT * FROM devices WHERE id = %s", (device_id,)) - device = cursor.fetchone() + new_device = cursor.fetchone() - # 设备创建成功后，在后台线程启动WEBRTC流处理 - threading.Thread( - target=run_webrtc_processing, - args=(device_data.ip, full_webrtc_url), - daemon=True # 设为守护线程，主程序退出时自动结束 - ).start() return APIResponse( code=200, - message="设备创建成功，已开始处理WEBRTC流", - data=DeviceResponse(**device) + message="设备创建成功", + data=DeviceResponse(**new_device) ) + except MySQLError as e: if conn: conn.rollback() raise Exception(f"创建设备失败：{str(e)}") from e except json.JSONDecodeError as e: - raise Exception(f"设备信息JSON序列化失败：{str(e)}") from e + raise Exception(f"设备详细信息JSON序列化失败：{str(e)}") from e except Exception as e: if conn: conn.rollback() @@ -148,139 +185,56 @@ async def create_device(request: Request, device_data: DeviceCreateRequest): db.close_connection(conn, cursor) -# ------------------------------ -# 2. 获取设备列表 -# ------------------------------ -@router.get("/", response_model=APIResponse, summary="获取设备列表") +@router.get("/", response_model=APIResponse, summary="获取设备列表（支持筛选分页）") async def get_device_list( - page: int = Query(1, ge=1, description="页码"), - page_size: int = Query(10, ge=1, le=100, description="每页条数"), - device_type: str = Query(None, description="设备类型筛选"), - online_status: int = Query(None, ge=0, le=1, description="在线状态筛选（1-在线、0-离线）") + page: int = Query(1, ge=1, description="页码，默认第1页"), + page_size: int = Query(10, ge=1, le=100, description="每页条数，1-100之间"), + device_type: str = Query(None, description="按设备类型筛选"), + online_status: int = Query(None, ge=0, le=1, description="按在线状态筛选") ): + # 原有代码保持不变 conn = None cursor = None try: conn = db.get_connection() cursor = conn.cursor(dictionary=True) - # 构建查询条件 where_clause = [] params = [] if device_type: where_clause.append("device_type = %s") params.append(device_type) - if online_status is not None: where_clause.append("device_online_status = %s") params.append(online_status) - # 总条数查询 - count_query = "SELECT COUNT(*) as total FROM devices" + count_query = "SELECT COUNT(*) AS total FROM devices" if where_clause: count_query += " WHERE " + " AND ".join(where_clause) - cursor.execute(count_query, params) total = cursor.fetchone()["total"] - # 分页查询（SELECT * 会自动匹配表字段、响应模型已对齐） offset = (page - 1) * page_size - query = "SELECT * FROM devices" + list_query = "SELECT * FROM devices" if where_clause: - query += " WHERE " + " AND ".join(where_clause) - query += " ORDER BY id DESC LIMIT %s OFFSET %s" + list_query += " WHERE " + " AND ".join(where_clause) + list_query += " ORDER BY id DESC LIMIT %s OFFSET %s" params.extend([page_size, offset]) - cursor.execute(query, params) - devices = cursor.fetchall() - - # 响应模型已更新为params字段、直接转换即可 - device_list = [DeviceResponse(**device) for device in devices] + cursor.execute(list_query, params) + device_list = cursor.fetchall() return APIResponse( code=200, message="获取设备列表成功", - data=DeviceListResponse(total=total, devices=device_list) + data=DeviceListResponse( + total=total, + devices=[DeviceResponse(**device) for device in device_list] + ) ) + except MySQLError as e: raise Exception(f"获取设备列表失败：{str(e)}") from e finally: db.close_connection(conn, cursor) - - -# ------------------------------ -# 3. 获取单个设备详情 -# ------------------------------ -@router.get("/{device_id}", response_model=APIResponse, summary="获取设备详情") -async def get_device_detail( - device_id: int, - current_user: UserResponse = Depends(get_current_user) -): - conn = None - cursor = None - try: - conn = db.get_connection() - cursor = conn.cursor(dictionary=True) - - # 查询设备信息（SELECT * 匹配表字段） - query = "SELECT * FROM devices WHERE id = %s" - cursor.execute(query, (device_id,)) - device = cursor.fetchone() - - if not device: - raise HTTPException( - status_code=404, - detail=f"设备ID为 {device_id} 的设备不存在" - ) - - # 响应模型已更新为params字段 - return APIResponse( - code=200, - message="获取设备详情成功", - data=DeviceResponse(**device) - ) - except MySQLError as e: - raise Exception(f"获取设备详情失败：{str(e)}") from e - finally: - db.close_connection(conn, cursor) - - -# ------------------------------ -# 4. 删除设备信息 -# ------------------------------ -@router.delete("/{device_id}", response_model=APIResponse, summary="删除设备信息") -async def delete_device( - device_id: int, - current_user: UserResponse = Depends(get_current_user) -): - conn = None - cursor = None - try: - conn = db.get_connection() - cursor = conn.cursor(dictionary=True) - - # 检查设备是否存在 - cursor.execute("SELECT id FROM devices WHERE id = %s", (device_id,)) - if not cursor.fetchone(): - raise HTTPException( - status_code=404, - detail=f"设备ID为 {device_id} 的设备不存在" - ) - - # 执行删除 - delete_query = "DELETE FROM devices WHERE id = %s" - cursor.execute(delete_query, (device_id,)) - conn.commit() - - return APIResponse( - code=200, - message=f"设备ID为 {device_id} 的设备已成功删除", - data=None - ) - except MySQLError as e: - if conn: - conn.rollback() - raise Exception(f"删除设备失败：{str(e)}") from e - finally: - db.close_connection(conn, cursor) diff --git a/ws/ws.py b/ws/ws.py index 8efca08..3e7e010 100644 --- a/ws/ws.py +++ b/ws/ws.py @@ -4,6 +4,9 @@ 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 @@ -11,7 +14,7 @@ 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" FORBIDDEN_WORDS_PATH = r"D:\Git\bin\video\ocr\forbidden_words.txt" OCR_CONFIG_PATH = r"D:\Git\bin\video\ocr\config\1.yaml" @@ -26,265 +29,287 @@ detector = MultiModelViolationDetector( ocr_confidence_threshold=0.5 ) -# -------------------------- 配置常量 -------------------------- +# 配置常量 HEARTBEAT_INTERVAL = 30 # 心跳检查间隔（秒） HEARTBEAT_TIMEOUT = 600 # 客户端超时阈值（秒） WS_ENDPOINT = "/ws" # WebSocket端点路径 -FRAME_QUEUE_SIZE = 1 # 帧队列大小限制（保持1，确保单帧处理） - -# -------------------------- 核心数据结构与全局变量 -------------------------- -ws_router = APIRouter() +FRAME_QUEUE_SIZE = 1 # 帧队列大小限制 -# 客户端连接封装（包含帧队列） +# 工具函数：获取格式化时间字符串（统一时间戳格式） +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) # 帧队列，长度为1 - self.consumer_task: Optional[asyncio.Task] = None # 消费者任务 + self.frame_queue = asyncio.Queue(maxsize=FRAME_QUEUE_SIZE) + self.consumer_task: Optional[asyncio.Task] = None - # 更新心跳时间 def update_heartbeat(self): + """更新心跳时间（客户端发送心跳时调用）""" self.last_heartbeat = datetime.datetime.now() - # 检查是否存活（超时返回False） 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 - # ---------- 新增：发送“允许发送二进制帧”的信号给客户端 ---------- - async def send_allow_send_frame(self): - """向客户端发送JSON信号，通知其可发送下一帧二进制数据""" + async def send_frame_permit(self): + """ + 发送「帧发送许可信号」 + 通知客户端可发送下一帧图像 + """ 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 # 可选：便于客户端确认自身身份 + frame_permit_msg = { + "type": "frame", + "timestamp": get_current_time_str(), + "client_ip": self.client_ip } - await self.websocket.send_json(allow_msg) - print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{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"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：发送「允许发送帧」信号失败 - {str(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() + + # -------------------------- 核心修改：取出帧后立即发送下一帧许可 -------------------------- + await self.send_frame_permit() # 取帧即通知客户端发下一帧，无需等处理完成 + # ----------------------------------------------------------------------------------------- + try: - # 处理帧数据 + # 2. 处理取出的帧（即使处理慢，客户端也已收到许可，可提前准备下一帧） await self.process_frame(frame_data) finally: - # 标记任务完成（队列计数-1，此时队列回到空状态） + # 3. 标记帧任务完成（无论处理成功/失败，都需清理队列） 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}：帧消费任务已取消") + print(f"[{get_current_time_str()}] 客户端{self.client_ip}：帧消费任务已取消") except Exception as e: - print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：帧处理错误 - {str(e)}") + print(f"[{get_current_time_str()}] 客户端{self.client_ip}：帧消费逻辑错误 - {str(e)}") async def process_frame(self, frame_data: bytes) -> None: - """处理单帧图像数据（原有逻辑不变）""" - # 将二进制数据转换为NumPy数组（uint8类型） + """处理单帧图像数据（检测违规后发送危险通知 + 调用违规次数加一方法）""" + # 二进制数据转OpenCV图像 nparr = np.frombuffer(frame_data, np.uint8) - # 解码为图像，返回与cv2.imread相同的格式（BGR通道的ndarray） img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) + if img is None: + print(f"[{get_current_time_str()}] 客户端{self.client_ip}：无法解析图像数据") + return - # 确保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/{self.client_ip.replace('.', '_')}_{timestamp}.jpg" + # 确保图像保存目录存在 + os.makedirs('images', exist_ok=True) + # 保存图像（按IP+时间戳命名，避免冲突） + filename = f"images/{self.client_ip.replace('.', '_')}_{get_current_time_file_str()}.jpg" try: - # 保存图像到本地 cv2.imwrite(filename, img) - print(f"[{datetime.datetime.now():%H:%M:%S}] 图像已保存至：{filename}") + print(f"[{get_current_time_str()}] 图像已保存至：{filename}") - # 进行检测 - if img is not None: - 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, - "violation_type": violation_type, - "details": details, - "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") - }) - else: - print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：未检测到任何违规内容") + # 执行违规检测 + has_violation, violation_type, details = 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"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：无法解析图像数据") + print(f"[{get_current_time_str()}] 客户端{self.client_ip}：未检测到违规") except Exception as e: - print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{self.client_ip}：图像处理错误 - {str(e)}") + print(f"[{get_current_time_str()}] 客户端{self.client_ip}：图像处理错误 - {str(e)}") -# 全局连接管理（IP -> 连接实例） +# 全局状态管理 connected_clients: Dict[str, ClientConnection] = {} -# 心跳任务（全局引用，用于关闭时清理） heartbeat_task: Optional[asyncio.Task] = None -# -------------------------- 心跳检查逻辑（原有逻辑不变） -------------------------- +# 心跳检查（定时清理超时客户端 + 调用离线状态更新方法） async def heartbeat_checker(): while True: - now = datetime.datetime.now() - # 1. 筛选超时客户端（避免遍历中修改字典） + current_time = get_current_time_str() timeout_ips = [ip for ip, conn in connected_clients.items() if not conn.is_alive()] - # 2. 处理超时连接（关闭+移除） if timeout_ips: - print(f"[{now:%H:%M:%S}] 心跳检查：{len(timeout_ips)}个客户端超时（{timeout_ips}）") + print(f"[{current_time}] 心跳检查：{len(timeout_ips)}个客户端超时（IP：{timeout_ips}）") for ip in timeout_ips: try: - # 取消消费者任务 - if connected_clients[ip].consumer_task and not connected_clients[ip].consumer_task.done(): - connected_clients[ip].consumer_task.cancel() - await connected_clients[ip].websocket.close(code=1008, reason="心跳超时") + 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="心跳超时") + + # 超时设为离线并记录 + 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: connected_clients.pop(ip, None) else: - print(f"[{now:%H:%M:%S}] 心跳检查：{len(connected_clients)}个客户端在线，无超时") + print(f"[{current_time}] 心跳检查：{len(connected_clients)}个客户端在线") - # 3. 等待下一轮检查 await asyncio.sleep(HEARTBEAT_INTERVAL) -# -------------------------- 应用生命周期（原有逻辑不变） -------------------------- +# 应用生命周期管理 @asynccontextmanager async def lifespan(app: FastAPI): global heartbeat_task - # 启动心跳任务 heartbeat_task = asyncio.create_task(heartbeat_checker()) - print(f"[{datetime.datetime.now():%H:%M:%S}] 心跳任务启动（ID：{id(heartbeat_task)}）") + 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"[{datetime.datetime.now():%H:%M:%S}] 心跳任务已取消") + print(f"[{get_current_time_str()}] 全局心跳检查任务已取消") except asyncio.CancelledError: pass -# -------------------------- 消息处理（文本/心跳逻辑不变，二进制逻辑保留） -------------------------- -async def send_heartbeat_ack(client_ip: str): - """回复心跳确认（原有逻辑不变）""" - if client_ip not in connected_clients: - return False +# 消息处理工具函数 +async def send_heartbeat_ack(conn: ClientConnection): try: - ack = { - "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), - "type": "heartbeat" + heartbeat_ack_msg = { + "type": "heart", + "timestamp": get_current_time_str(), + "client_ip": conn.client_ip } - await connected_clients[client_ip].websocket.send_json(ack) + await conn.websocket.send_json(heartbeat_ack_msg) + print(f"[{get_current_time_str()}] 客户端{conn.client_ip}：已发送心跳确认") return True - except Exception: - connected_clients.pop(client_ip, None) + except Exception as e: + 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(client_ip: str, text: str, conn: ClientConnection): - """处理文本消息（核心：心跳+JSON解析，原有逻辑不变）""" +async def handle_text_msg(conn: ClientConnection, text: str): try: msg = json.loads(text) - # 仅处理心跳类型消息 - if msg.get("type") == "heartbeat": + if msg.get("type") == "heart": conn.update_heartbeat() - await send_heartbeat_ack(client_ip) + await send_heartbeat_ack(conn) else: - print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：收到文本消息：{msg}") + print(f"[{get_current_time_str()}] 客户端{conn.client_ip}：未知文本消息类型（{msg.get('type')}）") except json.JSONDecodeError: - print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：无效JSON消息") + print(f"[{get_current_time_str()}] 客户端{conn.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 - - # 队列未满，添加帧到队列 +async def handle_binary_msg(conn: ClientConnection, data: bytes): try: conn.frame_queue.put_nowait(data) - print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：已接收{len(data)}字节二进制数据，加入队列") + print(f"[{get_current_time_str()}] 客户端{conn.client_ip}：图像数据（{len(data)}字节）已加入队列") except asyncio.QueueFull: - print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：队列突然满了，丢弃{len(data)}字节数据") + print(f"[{get_current_time_str()}] 客户端{conn.client_ip}：帧队列已满，丢弃当前图像数据") + + +# WebSocket路由配置 +ws_router = APIRouter() -# -------------------------- WebSocket核心端点（修改连接初始化逻辑） -------------------------- @ws_router.websocket(WS_ENDPOINT) async def websocket_endpoint(websocket: WebSocket): - # 接受连接 + 获取客户端IP await websocket.accept() - client_ip = websocket.client.host if websocket.client else "unknown" - now = datetime.datetime.now() - print(f"[{now:%H:%M:%S}] 客户端{client_ip}：连接成功") + 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 - consumer_task = None try: - # 处理重复连接（关闭旧连接） + # 处理重复连接 if client_ip in connected_clients: - # 取消旧连接的消费者任务 - if connected_clients[client_ip].consumer_task and not connected_clients[client_ip].consumer_task.done(): - connected_clients[client_ip].consumer_task.cancel() - await connected_clients[client_ip].websocket.close(code=1008, reason="同一IP新连接") + 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新连接建立") connected_clients.pop(client_ip) - print(f"[{now:%H:%M:%S}] 客户端{client_ip}：关闭旧连接") + print(f"[{current_time}] 客户端{client_ip}：已关闭旧连接") # 注册新连接 new_conn = ClientConnection(websocket, client_ip) connected_clients[client_ip] = new_conn + new_conn.start_consumer() + # 初始许可：连接建立后立即发一次，让客户端知道可发第一帧（后续靠取帧后自动发） + await new_conn.send_frame_permit() - # 启动帧消费任务 - consumer_task = new_conn.start_consumer() - # ---------- 修改：客户端刚连接时，队列空，立即发送「允许发送帧」信号 ---------- - await new_conn.send_allow_send_frame() - print(f"[{now:%H:%M:%S}] 客户端{client_ip}：注册成功，已启动帧消费任务，当前在线{len(connected_clients)}个") + # 标记上线并记录 + 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)}") - # 循环接收消息（原有逻辑不变） + print(f"[{current_time}] 客户端{client_ip}：新连接注册成功，在线数：{len(connected_clients)}") + + # 消息循环 while True: data = await websocket.receive() if "text" in data: - await handle_text_msg(client_ip, data["text"], new_conn) + await handle_text_msg(new_conn, data["text"]) elif "bytes" in data: - await handle_binary_msg(client_ip, data["bytes"]) + await handle_binary_msg(new_conn, data["bytes"]) - # 异常处理（断开/错误） except WebSocketDisconnect as e: - print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：主动断开（代码：{e.code}）") + print(f"[{get_current_time_str()}] 客户端{client_ip}：主动断开连接（代码：{e.code}）") except Exception as e: - print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：连接异常（{str(e)[:50]}）") + print(f"[{get_current_time_str()}] 客户端{client_ip}：连接异常 - {str(e)[:50]}") finally: - # 清理连接和任务 + # 清理资源并标记离线 if client_ip in connected_clients: - # 取消消费者任务 - if connected_clients[client_ip].consumer_task and not connected_clients[client_ip].consumer_task.done(): - connected_clients[client_ip].consumer_task.cancel() + 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) + 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) - print(f"[{datetime.datetime.now():%H:%M:%S}] 客户端{client_ip}：连接已清理，当前在线{len(connected_clients)}个") \ No newline at end of file + print(f"[{get_current_time_str()}] 客户端{client_ip}：资源已清理，在线数：{len(connected_clients)}")