From 4be7f7bf1495faefa06e89a3995a2ea3004ba35a Mon Sep 17 00:00:00 2001 From: ninghongbin <2409766686@qq.com> Date: Fri, 12 Sep 2025 14:05:09 +0800 Subject: [PATCH] =?UTF-8?q?=E7=9B=AE=E5=89=8D=E5=8F=AF=E4=BB=A5=E6=88=90?= =?UTF-8?q?=E5=8A=9F=E5=8A=A8=E6=80=81=E6=9B=B4=E6=8D=A2=E6=A8=A1=E5=9E=8B?= =?UTF-8?q?=E8=BF=90=E8=A1=8C=E7=9A=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app.py | 283 ++++++++++++++++++++++ core/all.py | 41 ++-- core/yolo.py | 47 ++-- main.py | 132 +++++++++-- schema/face_schema.py | 33 ++- schema/model_schema.py | 37 +++ schema/user_schema.py | 10 +- service/face_service.py | 364 +++++++++++++++------------- service/model_service.py | 497 +++++++++++++++++++++++++++++++++++++++ service/user_service.py | 106 ++++++++- util/face_util.py | 149 ++++++------ util/file_util.py | 83 +++++++ util/model_util.py | 61 +++++ 13 files changed, 1518 insertions(+), 325 deletions(-) create mode 100644 app.py create mode 100644 schema/model_schema.py create mode 100644 service/model_service.py create mode 100644 util/file_util.py create mode 100644 util/model_util.py diff --git a/app.py b/app.py new file mode 100644 index 0000000..59b48ed --- /dev/null +++ b/app.py @@ -0,0 +1,283 @@ +from flask import Flask, send_from_directory, abort, request +import os +import logging +from functools import wraps +from pathlib import Path +# 跨域依赖（必须安装：pip install flask-cors） +from flask_cors import CORS + +# 配置日志（保持原有格式） +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' +) +logger = logging.getLogger(__name__) + +# 初始化 Flask 应用（供 main.py 导入） +app = Flask(__name__) + +# ------------------------------ +# 核心修改：与 FastAPI 对齐的跨域配置 +# ------------------------------ +# 1. 允许的前端域名（完全复制 FastAPI 的 ALLOWED_ORIGINS，确保前后端一致） +ALLOWED_ORIGINS = [ + # "http://localhost:8080", # 本地前端开发地址（必改：替换为你的前端实际地址） + # "http://127.0.0.1:8080", + # "http://服务器IP:8080", # 部署后前端地址（替换为你的服务器IP/域名） + # # "*" 仅开发环境临时使用，生产环境必须删除（安全风险） + "*" +] + +# 2. 配置 CORS（与 FastAPI 规则完全对齐） +CORS( + app, + resources={ + r"/*": { # 对所有 Flask 路由生效（覆盖图片、模型下载所有接口） + "origins": ALLOWED_ORIGINS, # 允许的前端域名（与 FastAPI 一致） + "allow_credentials": True, # 允许携带 Cookie（与 FastAPI 一致，需登录态必开） + "methods": ["*"], # 允许所有 HTTP 方法（FastAPI 用 "*"，此处同步） + "allow_headers": ["*"], # 允许所有请求头（与 FastAPI 一致） + } + }, +) + +# ------------------------------ +# 核心路径配置（不变，确保资源目录正确） +# ------------------------------ +CURRENT_FILE_PATH = Path(__file__).resolve() +PROJECT_ROOT = CURRENT_FILE_PATH.parent # 项目根目录（video/） +# 资源目录（图片、模型） +BASE_IMAGE_DIR_DECT = str((PROJECT_ROOT / "resource" / "dect").resolve()) # 检测图片目录 +BASE_IMAGE_DIR_UP_IMAGES = str((PROJECT_ROOT / "up_images").resolve()) # 人脸图片目录 +BASE_MODEL_DIR = str((PROJECT_ROOT / "resource" / "models").resolve()) # 模型文件目录 + +# 打印路径配置（调试用，确认目录正确） +logger.info(f"[Flask 配置] 项目根目录：{PROJECT_ROOT}") +logger.info(f"[Flask 配置] 模型目录：{BASE_MODEL_DIR}") +logger.info(f"[Flask 配置] 人脸图片目录：{BASE_IMAGE_DIR_UP_IMAGES}") +logger.info(f"[Flask 配置] 检测图片目录：{BASE_IMAGE_DIR_DECT}") + +# ------------------------------ +# 安全检查装饰器（不变，防路径遍历/非法文件） +# ------------------------------ +def safe_path_check(root_dir: str): + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + resource_path = kwargs.get('resource_path', '').strip() + # 统一路径分隔符（兼容 Windows \ 和 Linux /） + resource_path = resource_path.replace("/", os.sep).replace("\\", os.sep) + # 拼接完整路径（防止路径遍历） + full_file_path = os.path.abspath(os.path.join(root_dir, resource_path)) + logger.debug( + f"[Flask 安全检查] 请求路径：{resource_path} | 完整路径：{full_file_path} | 根目录：{root_dir}" + ) + + # 1. 禁止路径遍历（确保请求文件在根目录内） + if not full_file_path.startswith(root_dir): + logger.warning( + f"[Flask 安全拦截] 非法路径遍历！IP：{request.remote_addr} | 请求路径：{resource_path}" + ) + abort(403) + + # 2. 检查文件存在且为有效文件（非目录） + if not os.path.exists(full_file_path) or not os.path.isfile(full_file_path): + logger.warning( + f"[Flask 资源错误] 文件不存在/非文件！IP：{request.remote_addr} | 路径：{full_file_path}" + ) + abort(404) + + # 3. 限制文件大小（模型200MB，图片10MB，避免超大文件攻击） + max_size = 200 * 1024 * 1024 if "models" in root_dir else 10 * 1024 * 1024 + if os.path.getsize(full_file_path) > max_size: + logger.warning( + f"[Flask 大小超限] 文件超过{max_size//1024//1024}MB！IP：{request.remote_addr} | 路径：{full_file_path}" + ) + abort(413) + + # 安全检查通过，传递根目录给视图函数 + return func(*args, **kwargs, root_dir=root_dir) + return wrapper + return decorator + +# ------------------------------ +# 1. 模型下载接口（/model/download/*） +# ------------------------------ +@app.route('/model/download/') +@safe_path_check(root_dir=BASE_MODEL_DIR) +def download_model(resource_path, root_dir): + try: + resource_path = resource_path.replace("/", os.sep).replace("\\", os.sep) + dir_path, file_name = os.path.split(resource_path) + full_dir = os.path.abspath(os.path.join(root_dir, dir_path)) + + # 仅允许 .pt 格式（YOLO 模型） + if not file_name.lower().endswith('.pt'): + logger.warning( + f"[Flask 格式错误] 非 .pt 模型文件！IP：{request.remote_addr} | 文件名：{file_name}" + ) + abort(415) + + logger.info( + f"[Flask 模型下载] 成功请求！IP：{request.remote_addr} | 文件：{file_name} | 目录：{full_dir}" + ) + + # 强制浏览器下载（而非预览），设置二进制文件类型 + return send_from_directory( + full_dir, + file_name, + as_attachment=True, + mimetype="application/octet-stream" + ) + + except Exception as e: + logger.error( + f"[Flask 模型下载异常] IP：{request.remote_addr} | 错误：{str(e)}" + ) + abort(500) + +# ------------------------------ +# 2. 人脸图片访问接口（/up_images/*） +# ------------------------------ +@app.route('/up_images/') +@safe_path_check(root_dir=BASE_IMAGE_DIR_UP_IMAGES) +def get_face_image(resource_path, root_dir): + try: + resource_path = resource_path.replace("/", os.sep).replace("\\", os.sep) + dir_path, file_name = os.path.split(resource_path) + full_dir = os.path.abspath(os.path.join(root_dir, dir_path)) + + # 仅允许常见图片格式 + allowed_ext = ('.png', '.jpg', '.jpeg', '.gif', '.bmp') + if not file_name.lower().endswith(allowed_ext): + logger.warning( + f"[Flask 格式错误] 非图片文件！IP：{request.remote_addr} | 文件名：{file_name}" + ) + abort(415) + + logger.info( + f"[Flask 人脸图片] 成功请求！IP：{request.remote_addr} | 文件：{file_name} | 目录：{full_dir}" + ) + + # 允许浏览器预览图片（而非下载） + return send_from_directory(full_dir, file_name, as_attachment=False) + + except Exception as e: + logger.error( + f"[Flask 人脸图片异常] IP：{request.remote_addr} | 错误：{str(e)}" + ) + abort(500) + +# ------------------------------ +# 3. 检测图片访问接口（/resource/dect/*） +# ------------------------------ +@app.route('/resource/dect/') +@safe_path_check(root_dir=BASE_IMAGE_DIR_DECT) +def get_dect_image(resource_path, root_dir): + try: + resource_path = resource_path.replace("/", os.sep).replace("\\", os.sep) + dir_path, file_name = os.path.split(resource_path) + full_dir = os.path.abspath(os.path.join(root_dir, dir_path)) + + # 仅允许常见图片格式 + allowed_ext = ('.png', '.jpg', '.jpeg', '.gif', '.bmp') + if not file_name.lower().endswith(allowed_ext): + logger.warning( + f"[Flask 格式错误] 非图片文件！IP：{request.remote_addr} | 文件名：{file_name}" + ) + abort(415) + + logger.info( + f"[Flask 检测图片] 成功请求！IP：{request.remote_addr} | 文件：{file_name} | 目录：{full_dir}" + ) + + return send_from_directory(full_dir, file_name, as_attachment=False) + + except Exception as e: + logger.error( + f"[Flask 检测图片异常] IP：{request.remote_addr} | 错误：{str(e)}" + ) + abort(500) + +# ------------------------------ +# 4. 兼容旧图片接口（/images/* → 映射到 /resource/dect/*） +# ------------------------------ +@app.route('/images/') +@safe_path_check(root_dir=BASE_IMAGE_DIR_DECT) +def get_compatible_image(resource_path, root_dir): + try: + # 逻辑与检测图片接口一致，仅URL前缀不同（兼容旧前端） + resource_path = resource_path.replace("/", os.sep).replace("\\", os.sep) + dir_path, file_name = os.path.split(resource_path) + full_dir = os.path.abspath(os.path.join(root_dir, dir_path)) + + allowed_ext = ('.png', '.jpg', '.jpeg', '.gif', '.bmp') + if not file_name.lower().endswith(allowed_ext): + logger.warning( + f"[Flask 格式错误] 非图片文件！IP：{request.remote_addr} | 文件名：{file_name}" + ) + abort(415) + + logger.info( + f"[Flask 兼容图片] 成功请求！IP：{request.remote_addr} | 文件：{file_name} | 目录：{full_dir}" + ) + + return send_from_directory(full_dir, file_name, as_attachment=False) + + except Exception as e: + logger.error( + f"[Flask 兼容图片异常] IP：{request.remote_addr} | 错误：{str(e)}" + ) + abort(500) + +# ------------------------------ +# 全局错误处理器（友好提示，与 FastAPI 错误信息风格一致） +# ------------------------------ +@app.errorhandler(403) +def forbidden_error(error): + return "❌ 禁止访问：路径非法（可能存在路径遍历）或无权限", 403 + +@app.errorhandler(404) +def not_found_error(error): + return "❌ 资源不存在：请检查URL路径（IP、目录、文件名）是否正确", 404 + +@app.errorhandler(413) +def too_large_error(error): + return "❌ 文件过大：图片最大10MB，模型最大200MB", 413 + +@app.errorhandler(415) +def unsupported_type_error(error): + return "❌ 不支持的文件类型：图片支持png/jpg/jpeg/gif/bmp，模型仅支持pt", 415 + +@app.errorhandler(500) +def server_error(error): + return "❌ 服务器内部错误：请联系管理员查看后台日志", 500 + +# ------------------------------ +# Flask 独立启动入口（供测试，实际由 main.py 子线程启动） +# ------------------------------ +if __name__ == '__main__': + # 确保所有资源目录存在（防止初始化失败） + required_dirs = [ + (BASE_IMAGE_DIR_DECT, "检测图片目录"), + (BASE_IMAGE_DIR_UP_IMAGES, "人脸图片目录"), + (BASE_MODEL_DIR, "模型文件目录") + ] + for dir_path, dir_desc in required_dirs: + if not os.path.exists(dir_path): + logger.info(f"[Flask 初始化] {dir_desc}不存在，创建：{dir_path}") + os.makedirs(dir_path, exist_ok=True) + + # 启动提示（含访问示例） + logger.info("\n[Flask 服务启动成功！] 支持的接口：") + logger.info(f"1. 模型下载 → http://服务器IP:5000/model/download/resource/models/xxx.pt") + logger.info(f"2. 人脸图片 → http://服务器IP:5000/up_images/xxx.jpg") + logger.info(f"3. 检测图片 → http://服务器IP:5000/resource/dect/xxx.jpg 或 http://服务器IP:5000/images/xxx.jpg\n") + + # 启动服务（禁用 debug 和自动重载，避免多线程冲突） + app.run( + host="0.0.0.0", # 允许外部IP访问 + port=5000, # 与 main.py 中 Flask 端口一致 + debug=False, + use_reloader=False + ) \ No newline at end of file diff --git a/core/all.py b/core/all.py index 7cc307d..2e36f91 100644 --- a/core/all.py +++ b/core/all.py @@ -57,50 +57,39 @@ def save_db(model_type, client_ip, result): +# 修正后的 detect 函数关键部分 def detect(client_ip, frame): - """ - 执行模型检测，检测到违规时按指定格式保存图片 - 参数： - frame: 待检测的图像帧（OpenCV格式，numpy.ndarray类型） - 返回： - (检测结果布尔值, 检测详情, 检测模型类型) - """ - # 1. YOLO检测（优先级1） + # 1. YOLO检测 yolo_flag, yolo_result = yoloDetect(frame) - print(f"YOLO检测结果：{yolo_result}") if yolo_flag: + # model_type 传入 "yolo"（正确） full_save_path, display_path = get_image_save_path(model_type="yolo", client_ip=client_ip) - if full_save_path: # 只判断完整路径是否有效（用于保存） + if full_save_path: cv2.imwrite(full_save_path, frame) - # 打印时使用「显示用短路径」，符合需求格式 - print(f"✅ YOLO违规图片已保存：{display_path}") + print(f"✅ yolo违规图片已保存：{display_path}") # 日志也修正 save_db(model_type="yolo", client_ip=client_ip, result=str(full_save_path)) return (True, yolo_result, "yolo") - # - # # 2. 人脸检测（优先级2） + + # 2. 人脸检测 face_flag, face_result = faceDetect(frame) - print(f"人脸检测结果：{face_result}") if face_flag: - full_save_path, display_path = get_image_save_path(model_type="yolo", client_ip=client_ip) - if full_save_path: # 只判断完整路径是否有效（用于保存） + full_save_path, display_path = get_image_save_path(model_type="face", client_ip=client_ip) # 这里改了 + if full_save_path: cv2.imwrite(full_save_path, frame) - # 打印时使用「显示用短路径」，符合需求格式 - print(f"✅ face违规图片已保存：{display_path}") + print(f"✅ face违规图片已保存：{display_path}") # 日志也修正 save_db(model_type="face", client_ip=client_ip, result=str(full_save_path)) return (True, face_result, "face") - # 3. OCR检测（优先级3） + # 3. OCR检测 ocr_flag, ocr_result = ocrDetect(frame) - print(f"OCR检测结果：{ocr_result}") if ocr_flag: - # 解构元组，保存用完整路径，打印用短路径 - full_save_path, display_path = get_image_save_path(model_type="yolo", client_ip=client_ip) - if full_save_path: # 只判断完整路径是否有效（用于保存） + full_save_path, display_path = get_image_save_path(model_type="ocr", client_ip=client_ip) # 这里改了 + if full_save_path: cv2.imwrite(full_save_path, frame) - # 打印时使用「显示用短路径」，符合需求格式 - print(f"✅ ocr违规图片已保存：{display_path}") + print(f"✅ ocr违规图片已保存：{display_path}") # 日志也修正 save_db(model_type="ocr", client_ip=client_ip, result=str(full_save_path)) return (True, ocr_result, "ocr") + # 4. 无违规内容（不保存图片） print(f"❌ 未检测到任何违规内容，不保存图片") return (False, "未检测到任何内容", "none") \ No newline at end of file diff --git a/core/yolo.py b/core/yolo.py index 54d9c2b..d7c299e 100644 --- a/core/yolo.py +++ b/core/yolo.py @@ -1,37 +1,43 @@ import os - +import numpy as np from ultralytics import YOLO +from service.model_service import get_current_yolo_model # 从模型管理模块获取模型 -# 全局变量 +# 全局模型变量 _yolo_model = None -model_path = os.path.join(os.path.dirname(__file__), "models", "best.pt") - - -def load_model(): - """加载YOLO目标检测模型""" +def load_model(model_path=None): + """加载YOLO模型（优先使用模型管理模块的默认模型）""" global _yolo_model + if model_path is None: + _yolo_model = get_current_yolo_model() + return _yolo_model is not None + try: _yolo_model = YOLO(model_path) + return True except Exception as e: - print(f"YOLO model load failed: {e}") + print(f"YOLO模型加载失败（指定路径）：{str(e)}") return False - return True if _yolo_model else False - def detect(frame, conf_threshold=0.2): - """YOLO目标检测、返回(是否识别到, 结果字符串)""" + """执行目标检测，返回(是否成功, 结果字符串)""" global _yolo_model - if not _yolo_model or frame is None: - return (False, "未初始化或无效帧") + # 确保模型已加载 + if not _yolo_model: + if not load_model(): + return (False, "模型未初始化") + + if frame is None: + return (False, "无效输入帧") try: - results = _yolo_model(frame, conf=conf_threshold) - # 检查是否有检测结果 + # 执行检测（frame应为numpy数组） + results = _yolo_model(frame, conf=conf_threshold, verbose=False) has_results = len(results[0].boxes) > 0 if results else False if not has_results: @@ -42,13 +48,12 @@ def detect(frame, conf_threshold=0.2): for box in results[0].boxes: cls = int(box.cls[0]) conf = float(box.conf[0]) - bbox = [float(x) for x in box.xyxy[0]] + bbox = [round(x, 2) for x in box.xyxy[0].tolist()] # 保留两位小数 class_name = _yolo_model.names[cls] if hasattr(_yolo_model, 'names') else f"类别{cls}" - result_parts.append(f"{class_name} (置信度: {conf:.2f}, 边界框: {bbox})") + result_parts.append(f"{class_name}（置信度：{conf:.2f}，位置：{bbox}）") - result_str = "; ".join(result_parts) - return (has_results, result_str) + return (True, "; ".join(result_parts)) except Exception as e: - print(f"YOLO detect error: {e}") - return (False, f"检测错误: {str(e)}") \ No newline at end of file + print(f"检测过程出错：{str(e)}") + return (False, f"检测错误：{str(e)}") diff --git a/main.py b/main.py index 271c78d..84205ef 100644 --- a/main.py +++ b/main.py @@ -1,62 +1,142 @@ -from PIL import Image # 正确导入 -import numpy as np - -import uvicorn from PIL import Image +import numpy as np +import uvicorn +import threading +import time +import os from fastapi import FastAPI -from core.all import load_model,detect +# 新增：导入 CORS 相关依赖 +from fastapi.middleware.cors import CORSMiddleware + +# 导入 Flask 服务实例 +from app import app as flask_app + +# 原有业务导入 +from core.all import load_model, detect from ds.config import SERVER_CONFIG from middle.error_handler import global_exception_handler from service.user_service import router as user_router from service.sensitive_service import router as sensitive_router from service.face_service import router as face_router from service.device_service import router as device_router +from service.model_service import router as model_router # 模型管理路由 from ws.ws import ws_router, lifespan from core.establish import create_directory_structure -# ------------------------------ -# 初始化 FastAPI 应用、指定生命周期管理 -# ------------------------------ + +# Flask 服务启动函数（不变） +def start_flask_service(): + try: + print(f"\n[Flask 服务] 准备启动，端口：5000") + print(f"[Flask 服务] 访问示例：http://服务器IP:5000/resource/dect/ocr/xxx.jpg\n") + + BASE_IMAGE_DIR = os.path.abspath(os.path.join("resource", "dect")) + if not os.path.exists(BASE_IMAGE_DIR): + print(f"[Flask 服务] 图片根目录不存在，创建：{BASE_IMAGE_DIR}") + os.makedirs(BASE_IMAGE_DIR, exist_ok=True) + + flask_app.run( + host="0.0.0.0", + port=5000, + debug=False, + use_reloader=False + ) + except Exception as e: + print(f"[Flask 服务] 启动失败：{str(e)}") + + +# 初始化 FastAPI 应用（新增 CORS 配置） app = FastAPI( title="内容安全审核后台", - description="内容安全审核后台", + description="含图片访问服务和动态模型管理", version="1.0.0", lifespan=lifespan ) # ------------------------------ -# 注册路由 +# 新增：完整 CORS 配置（解决跨域问题） # ------------------------------ +# 1. 允许的前端域名（根据实际情况修改！本地开发通常是 http://localhost:8080 等） +ALLOWED_ORIGINS = [ + # "http://localhost:8080", # 前端本地开发地址（必改，填实际前端地址） + # "http://127.0.0.1:8080", + # "http://服务器IP:8080", # 部署后前端地址（如适用） + "*" #表示允许所有域名（开发环境可用，生产环境不推荐） +] + +# 2. 配置 CORS 中间件 +app.add_middleware( + CORSMiddleware, + allow_origins=ALLOWED_ORIGINS, # 允许的前端域名 + allow_credentials=True, # 允许携带 Cookie（如需登录态则必开） + allow_methods=["*"], # 允许所有 HTTP 方法（包括 PUT/DELETE） + allow_headers=["*"], # 允许所有请求头（包括 Content-Type） +) + +# 注册路由（不变） app.include_router(user_router) app.include_router(device_router) app.include_router(face_router) app.include_router(sensitive_router) +app.include_router(model_router) # 模型管理路由 app.include_router(ws_router) -# ------------------------------ -# 注册全局异常处理器 -# ------------------------------ +# 注册全局异常处理器（不变） app.add_exception_handler(Exception, global_exception_handler) -# ------------------------------ -# 启动服务 -# ------------------------------ +# 主服务启动入口（不变） if __name__ == "__main__": - # -------------------------- 配置调整 -------------------------- - # 模型配置路径（建议改为环境变量） - YOLO_MODEL_PATH = r"/core/models\best.pt" - OCR_CONFIG_PATH = r"/core/config\config.yaml" - + # 1. 初始化资源 create_directory_structure() + print(f"[初始化] 目录结构创建完成") - # 初始化项目（默认端口设为8000、避免初始化失败时port未定义） + # 创建模型保存目录 + MODEL_SAVE_DIR = os.path.join("core", "models") + os.makedirs(MODEL_SAVE_DIR, exist_ok=True) + print(f"[初始化] 模型保存目录：{MODEL_SAVE_DIR}") + + # # 模型路径配置 + # YOLO_MODEL_PATH = os.path.join("core", "models", "best.pt") + # OCR_CONFIG_PATH = os.path.join("core", "config", "config.yaml") + # print(f"[初始化] 默认YOLO模型路径：{YOLO_MODEL_PATH}") + # print(f"[初始化] OCR 配置路径：{OCR_CONFIG_PATH}") + + # 加载检测模型 + try: + load_success = load_model() + if load_success: + print(f"[初始化] 检测模型加载完成") + else: + print(f"[初始化] 未找到默认模型，可通过API上传并设置默认模型") + except Exception as e: + print(f"[初始化] 模型加载警告：{str(e)}（服务仍启动）") + + + + # 2. 启动 Flask 服务（子线程） + flask_thread = threading.Thread( + target=start_flask_service, + daemon=True + ) + flask_thread.start() + + # 等待 Flask 初始化 + time.sleep(1) + if flask_thread.is_alive(): + print(f"[Flask 服务] 启动成功（运行中）") + else: + print(f"[Flask 服务] 启动失败！图片访问不可用") + + # 3. 启动 FastAPI 主服务 port = int(SERVER_CONFIG.get("port", 8000)) + print(f"\n[FastAPI 服务] 准备启动，端口：{port}") + print(f"[FastAPI 服务] 接口文档：http://服务器IP:{port}/docs\n") - # 启动 UVicorn 服务 uvicorn.run( app="main:app", host="0.0.0.0", port=port, - workers=8, - ws="websockets" - ) + workers=1, + ws="websockets", + reload=False + ) \ No newline at end of file diff --git a/schema/face_schema.py b/schema/face_schema.py index 4ade2b9..8e2663e 100644 --- a/schema/face_schema.py +++ b/schema/face_schema.py @@ -1,30 +1,41 @@ from datetime import datetime from pydantic import BaseModel, Field +from typing import List, Optional # ------------------------------ -# 请求模型（前端传参校验） +# 请求模型（前端传参校验）- 保留update的eigenvalue（如需更新特征值） # ------------------------------ class FaceCreateRequest(BaseModel): """创建人脸记录请求模型（无需ID、由数据库自增）""" - name: str = Field(None, max_length=255, description="名称（可选、最长255字符）") + name: Optional[str] = Field(None, max_length=255, description="名称（可选、最长255字符）") class FaceUpdateRequest(BaseModel): - """更新人脸记录请求模型（不变）""" - name: str = Field(None, max_length=255, description="名称") - eigenvalue: str = Field(None, max_length=255, description="特征（文件处理后可更新）") + """更新人脸记录请求模型 - 保留eigenvalue（如需更新特征值，不影响返回）""" + name: Optional[str] = Field(None, max_length=255, description="名称（可选）") + eigenvalue: Optional[str] = Field(None, description="特征值（可选，文件处理后可更新）") # 保留更新能力 + address: Optional[str] = Field(None, description="图片完整路径（可选，更新图片时使用）") # ------------------------------ -# 响应模型（后端返回数据） +# 响应模型（后端返回数据）- 核心修改：删除eigenvalue字段 # ------------------------------ class FaceResponse(BaseModel): - """人脸记录响应模型（仍包含ID、由数据库生成后返回）""" + """人脸记录响应模型（仅返回需要的字段，移除eigenvalue）""" id: int = Field(..., description="主键ID（数据库自增）") - name: str = Field(None, description="名称") - eigenvalue: str | None = Field(None, description="特征（可为空）") - created_at: datetime = Field(..., description="记录创建时间") - updated_at: datetime = Field(..., description="记录更新时间") + name: Optional[str] = Field(None, description="名称") + address: Optional[str] = Field(None, description="人脸图片完整保存路径（数据库新增字段）") # 仅保留address + created_at: datetime = Field(..., description="记录创建时间（数据库自动生成）") + updated_at: datetime = Field(..., description="记录更新时间（数据库自动生成）") + # 关键配置：支持从数据库查询结果（字典）直接转换 model_config = {"from_attributes": True} + + +class FaceListResponse(BaseModel): + """人脸列表分页响应模型（结构不变，内部FaceResponse已移除eigenvalue）""" + total: int = Field(..., description="筛选后的总记录数") + faces: List[FaceResponse] = Field(..., description="当前页的人脸记录列表") + + model_config = {"from_attributes": True} \ No newline at end of file diff --git a/schema/model_schema.py b/schema/model_schema.py new file mode 100644 index 0000000..8c4ebcc --- /dev/null +++ b/schema/model_schema.py @@ -0,0 +1,37 @@ +from datetime import datetime +from pydantic import BaseModel, Field +from typing import List, Optional + + +# 请求模型 +class ModelCreateRequest(BaseModel): + name: str = Field(..., max_length=255, description="模型名称（必填，如：yolo-v8s-car）") + description: Optional[str] = Field(None, description="模型描述（可选）") + is_default: Optional[bool] = Field(False, description="是否设为默认模型") + + +class ModelUpdateRequest(BaseModel): + name: Optional[str] = Field(None, max_length=255, description="模型名称（可选修改）") + description: Optional[str] = Field(None, description="模型描述（可选修改）") + is_default: Optional[bool] = Field(None, description="是否设为默认模型（可选切换）") + + +# 响应模型 +class ModelResponse(BaseModel): + id: int = Field(..., description="模型ID") + name: str = Field(..., description="模型名称") + path: str = Field(..., description="模型文件相对路径") + is_default: bool = Field(..., description="是否默认模型") + description: Optional[str] = Field(None, description="模型描述") + file_size: Optional[int] = Field(None, description="文件大小（字节）") + created_at: datetime = Field(..., description="创建时间") + updated_at: datetime = Field(..., description="更新时间") + + model_config = {"from_attributes": True} + + +class ModelListResponse(BaseModel): + total: int = Field(..., description="总记录数") + models: List[ModelResponse] = Field(..., description="当前页模型列表") + + model_config = {"from_attributes": True} diff --git a/schema/user_schema.py b/schema/user_schema.py index 6d8d9b1..30a86cc 100644 --- a/schema/user_schema.py +++ b/schema/user_schema.py @@ -1,6 +1,6 @@ from datetime import datetime - from pydantic import BaseModel, Field +from typing import List, Optional # ------------------------------ @@ -30,3 +30,11 @@ class UserResponse(BaseModel): # Pydantic V2 配置（支持从数据库查询结果转换） model_config = {"from_attributes": True} + + +class UserListResponse(BaseModel): + """用户列表分页响应模型（与设备/人脸列表结构对齐）""" + total: int = Field(..., description="用户总数") + users: List[UserResponse] = Field(..., description="当前页用户列表") + + model_config = {"from_attributes": True} \ No newline at end of file diff --git a/service/face_service.py b/service/face_service.py index 544a45a..0a6c8ad 100644 --- a/service/face_service.py +++ b/service/face_service.py @@ -1,162 +1,140 @@ -from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form +from fastapi import APIRouter, HTTPException, UploadFile, File, Form, Query, Request +from fastapi.responses import FileResponse from mysql.connector import Error as MySQLError +import os +from pathlib import Path from ds.db import db -from schema.face_schema import FaceCreateRequest, FaceUpdateRequest, FaceResponse -from schema.response_schema import APIResponse -from middle.auth_middleware import get_current_user -from schema.user_schema import UserResponse - -from util.face_util import add_binary_data,get_average_feature -#初始化实例 - -router = APIRouter( - prefix="/faces", - tags=["人脸管理"] +from schema.face_schema import ( + FaceCreateRequest, + FaceUpdateRequest, + FaceResponse, + FaceListResponse ) +from schema.response_schema import APIResponse +from util.face_util import add_binary_data, get_average_feature +from util.file_util import save_face_to_up_images + +router = APIRouter(prefix="/faces", tags=["人脸管理"]) # ------------------------------ -# 1. 创建人脸记录（核心修正: ID 数据库自增、前端无需传） +# 1. 创建人脸记录（使用修复后的路径） # ------------------------------ -@router.post("", response_model=APIResponse, summary="创建人脸记录（传名称+文件、ID自增）") +@router.post("", response_model=APIResponse, summary="创建人脸记录") async def create_face( - # 前端仅需传: name（可选、Form格式）、file（必传、文件） + request: Request, name: str = Form(None, max_length=255, description="名称（可选）"), - file: UploadFile = File(..., description="人脸文件（必传、暂不处理内容）") + file: UploadFile = File(..., description="人脸文件（必传）") ): - """ - 创建人脸记录: - - 需登录认证 - - 前端传参: multipart/form-data 表单（name 可选、file 必传） - - ID 由数据库自动生成、无需前端传入 - - 暂不处理文件内容、eigenvalue 设为 None - """ - - # 调用你的方法 conn = None cursor = None try: - # 1. 用模型校验 name（仅校验长度、无需ID） face_create = FaceCreateRequest(name=name) + client_ip = request.client.host if request.client else "" + if not client_ip: + raise HTTPException(status_code=400, detail="无法获取客户端IP") conn = db.get_connection() cursor = conn.cursor(dictionary=True) - # 把文件转为二进制数组 + # 读取图片并保存（使用修复后的路径逻辑） file_content = await file.read() + file_ext = file.filename.split(".")[-1].lower() if "." in file.filename else "jpg" + save_result = save_face_to_up_images( + client_ip=client_ip, + face_name=name, + image_bytes=file_content, + image_format=file_ext + ) + if not save_result["success"]: + raise HTTPException(status_code=500, detail=f"图片保存失败：{save_result['msg']}") + db_image_path = save_result["db_path"] # 从修复后的方法获取路径 - # 计算特征值 - flag, eigenvalue = add_binary_data(file_content) + # 提取人脸特征 + detect_success, detect_result = add_binary_data(file_content) + if not detect_success: + raise HTTPException(status_code=400, detail=f"人脸检测失败：{detect_result}") + eigenvalue = detect_result - if flag == False: - raise HTTPException( - status_code=500, - detail="未检测到人脸" - ) - - # 打印数组长度 - print(f"文件大小: {len(file_content)} 字节") - - # 2. 插入数据库: 无需传 ID（自增）、只传 name 和 eigenvalue（None） + # 插入数据库 insert_query = """ - INSERT INTO face (name, eigenvalue) - VALUES (%s, %s) + INSERT INTO face (name, eigenvalue, address) + VALUES (%s, %s, %s) """ - cursor.execute(insert_query, (face_create.name, str(eigenvalue))) + cursor.execute(insert_query, (face_create.name, str(eigenvalue), db_image_path)) conn.commit() - # 3. 获取数据库自动生成的 ID（关键: 用 LAST_INSERT_ID() 查刚插入的记录） - select_new_query = "SELECT * FROM face WHERE id = LAST_INSERT_ID()" - cursor.execute(select_new_query) + # 查询新记录 + cursor.execute(""" + SELECT id, name, address, created_at, updated_at + FROM face + WHERE id = LAST_INSERT_ID() + """) created_face = cursor.fetchone() - if not created_face: - raise HTTPException( - status_code=500, - detail="创建人脸记录成功、但无法获取新创建的记录" - ) + raise HTTPException(status_code=500, detail="创建成功但无法获取记录") return APIResponse( code=201, - message=f"人脸记录创建成功（ID: {created_face['id']}、文件名: {file.filename}）", - data=FaceResponse(** created_face) + message=f"人脸记录创建成功（ID: {created_face['id']}）", + data=FaceResponse(**created_face) ) except MySQLError as e: if conn: conn.rollback() - # 改为使用HTTPException - raise HTTPException( - status_code=500, - detail=f"创建人脸记录失败: {str(e)}" - ) from e + raise HTTPException(status_code=500, detail=f"创建失败: {str(e)}") from e except Exception as e: - # 捕获其他可能的异常 - raise HTTPException( - status_code=500, - detail=f"服务器错误: {str(e)}" - ) from e + raise HTTPException(status_code=500, detail=f"服务器错误: {str(e)}") from e finally: - await file.close() # 关闭文件流 + await file.close() db.close_connection(conn, cursor) - # 调用人脸识别得到特征值（这里可以添加你的人脸识别逻辑） - flag, eigenvalue = add_binary_data(file_content) - if flag == False: - raise HTTPException( - status_code=500, - detail="未检测到人脸" - ) - - # 将 eigenvalue 转为 str - eigenvalue = str(eigenvalue) +# 其他接口（获取单条/列表、更新、删除、获取图片）与之前一致，无需修改 # ------------------------------ -# 2. 获取单个人脸记录（不变、用自增ID查询） +# 2. 获取单个人脸记录 # ------------------------------ @router.get("/{face_id}", response_model=APIResponse, summary="获取单个人脸记录") -async def get_face( - face_id: int, # 这里的 ID 是数据库自增的、前端从创建响应中获取 - current_user: UserResponse = Depends(get_current_user) -): +async def get_face(face_id: int): conn = None cursor = None try: conn = db.get_connection() cursor = conn.cursor(dictionary=True) - query = "SELECT * FROM face WHERE id = %s" + query = """ + SELECT id, name, address, created_at, updated_at + FROM face + WHERE id = %s + """ cursor.execute(query, (face_id,)) face = cursor.fetchone() if not face: - raise HTTPException( - status_code=404, - detail=f"ID为 {face_id} 的人脸记录不存在" - ) + raise HTTPException(status_code=404, detail=f"ID为 {face_id} 的记录不存在") return APIResponse( code=200, - message="人脸记录查询成功", + message="查询成功", data=FaceResponse(**face) ) except MySQLError as e: - # 改为使用HTTPException - raise HTTPException( - status_code=500, - detail=f"查询人脸记录失败: {str(e)}" - ) from e + raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}") from e finally: db.close_connection(conn, cursor) -# 后续 3.获取所有、4.更新、5.删除 接口（修复异常处理） # ------------------------------ -# 3. 获取所有人脸记录（不变） +# 3. 获取人脸列表 # ------------------------------ -@router.get("", response_model=APIResponse, summary="获取所有人脸记录") -async def get_all_faces( +@router.get("", response_model=APIResponse, summary="获取人脸列表（分页+筛选）") +async def get_face_list( + page: int = Query(1, ge=1), + page_size: int = Query(10, ge=1, le=100), + name: str = Query(None), + has_eigenvalue: bool = Query(None) ): conn = None cursor = None @@ -164,50 +142,66 @@ async def get_all_faces( conn = db.get_connection() cursor = conn.cursor(dictionary=True) - query = "SELECT * FROM face ORDER BY id" # 按自增ID排序 - cursor.execute(query) - faces = cursor.fetchall() + where_clause = [] + params = [] + if name: + where_clause.append("name LIKE %s") + params.append(f"%{name}%") + if has_eigenvalue is not None: + where_clause.append("eigenvalue IS NOT NULL" if has_eigenvalue else "eigenvalue IS NULL") + + # 总记录数 + count_query = "SELECT COUNT(*) AS total FROM face" + if where_clause: + count_query += " WHERE " + " AND ".join(where_clause) + cursor.execute(count_query, params) + total = cursor.fetchone()["total"] + + # 列表数据 + offset = (page - 1) * page_size + list_query = """ + SELECT id, name, address, created_at, updated_at + FROM face + """ + if where_clause: + list_query += " WHERE " + " AND ".join(where_clause) + list_query += " ORDER BY id DESC LIMIT %s OFFSET %s" + params.extend([page_size, offset]) + + cursor.execute(list_query, params) + face_list = cursor.fetchall() return APIResponse( code=200, - message="所有人脸记录查询成功", - data=[FaceResponse(** face) for face in faces] + message=f"获取成功（共{total}条）", + data=FaceListResponse( + total=total, + faces=[FaceResponse(**face) for face in face_list] + ) ) except MySQLError as e: - raise HTTPException( - status_code=500, - detail=f"查询所有人脸记录失败: {str(e)}" - ) from e + raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}") from e finally: db.close_connection(conn, cursor) # ------------------------------ -# 4. 更新人脸记录（不变、用自增ID更新） +# 4. 更新人脸记录 # ------------------------------ @router.put("/{face_id}", response_model=APIResponse, summary="更新人脸记录") -async def update_face( - face_id: int, - face_update: FaceUpdateRequest, - current_user: UserResponse = Depends(get_current_user) -): +async def update_face(face_id: int, face_update: FaceUpdateRequest): conn = None cursor = None try: conn = db.get_connection() cursor = conn.cursor(dictionary=True) - # 检查记录是否存在 - check_query = "SELECT id FROM face WHERE id = %s" - cursor.execute(check_query, (face_id,)) - existing_face = cursor.fetchone() - if not existing_face: - raise HTTPException( - status_code=404, - detail=f"ID为 {face_id} 的人脸记录不存在" - ) + cursor.execute("SELECT id, address FROM face WHERE id = %s", (face_id,)) + exist_face = cursor.fetchone() + if not exist_face: + raise HTTPException(status_code=404, detail=f"ID为 {face_id} 的记录不存在") + old_db_path = exist_face["address"] - # 构建更新语句 update_fields = [] params = [] if face_update.name is not None: @@ -216,6 +210,18 @@ async def update_face( if face_update.eigenvalue is not None: update_fields.append("eigenvalue = %s") params.append(face_update.eigenvalue) + if face_update.address is not None: + # 删除旧图片（相对路径转绝对路径） + if old_db_path: + old_abs_path = Path(old_db_path).resolve() + if old_abs_path.exists(): + try: + old_abs_path.unlink() # 使用Path方法删除更安全 + print(f"[FaceRouter] 已删除旧图片：{old_abs_path}") + except Exception as e: + print(f"[FaceRouter] 删除旧图片失败：{str(e)}") + update_fields.append("address = %s") + params.append(face_update.address) if not update_fields: raise HTTPException(status_code=400, detail="至少需提供一个更新字段") @@ -225,117 +231,143 @@ async def update_face( cursor.execute(update_query, params) conn.commit() - # 查询更新后记录 - select_query = "SELECT * FROM face WHERE id = %s" - cursor.execute(select_query, (face_id,)) + cursor.execute(""" + SELECT id, name, address, created_at, updated_at + FROM face + WHERE id = %s + """, (face_id,)) updated_face = cursor.fetchone() return APIResponse( code=200, - message="人脸记录更新成功", + message="更新成功", data=FaceResponse(**updated_face) ) except MySQLError as e: if conn: conn.rollback() - raise HTTPException( - status_code=500, - detail=f"更新人脸记录失败: {str(e)}" - ) from e + raise HTTPException(status_code=500, detail=f"更新失败: {str(e)}") from e finally: db.close_connection(conn, cursor) # ------------------------------ -# 5. 删除人脸记录（不变、用自增ID删除） +# 5. 删除人脸记录 # ------------------------------ @router.delete("/{face_id}", response_model=APIResponse, summary="删除人脸记录") -async def delete_face( - face_id: int, - current_user: UserResponse = Depends(get_current_user) -): +async def delete_face(face_id: int): conn = None cursor = None try: conn = db.get_connection() cursor = conn.cursor(dictionary=True) - check_query = "SELECT id FROM face WHERE id = %s" - cursor.execute(check_query, (face_id,)) - existing_face = cursor.fetchone() - if not existing_face: - raise HTTPException( - status_code=404, - detail=f"ID为 {face_id} 的人脸记录不存在" - ) + cursor.execute("SELECT id, address FROM face WHERE id = %s", (face_id,)) + exist_face = cursor.fetchone() + if not exist_face: + raise HTTPException(status_code=404, detail=f"ID为 {face_id} 的记录不存在") + old_db_path = exist_face["address"] - delete_query = "DELETE FROM face WHERE id = %s" - cursor.execute(delete_query, (face_id,)) + cursor.execute("DELETE FROM face WHERE id = %s", (face_id,)) conn.commit() + # 删除图片 + extra_msg = "" + if old_db_path: + old_abs_path = Path(old_db_path).resolve() + if old_abs_path.exists(): + try: + old_abs_path.unlink() + print(f"[FaceRouter] 已删除图片：{old_abs_path}") + extra_msg = "（已同步删除图片）" + except Exception as e: + print(f"[FaceRouter] 删除图片失败：{str(e)}") + extra_msg = "（图片删除失败）" + else: + extra_msg = "（图片不存在）" + else: + extra_msg = "（无关联图片）" + return APIResponse( code=200, - message=f"ID为 {face_id} 的人脸记录删除成功", + message=f"ID为 {face_id} 的记录删除成功 {extra_msg}", data=None ) except MySQLError as e: if conn: conn.rollback() - raise HTTPException( - status_code=500, - detail=f"删除人脸记录失败: {str(e)}" - ) from e + raise HTTPException(status_code=500, detail=f"删除失败: {str(e)}") from e finally: db.close_connection(conn, cursor) -def get_all_face_name_with_eigenvalue() -> dict: - """ - 获取所有人脸的名称及其对应的特征值、组成字典返回 - key: 人脸名称（name） - value: 人脸特征值（eigenvalue）、若名称重复则返回平均特征值 - 注: 过滤掉name为None的记录、避免字典key为None的情况 - """ +# ------------------------------ +# 6. 获取人脸图片 +# ------------------------------ +@router.get("/{face_id}/image", summary="获取人脸图片") +async def get_face_image(face_id: int): + conn = None + cursor = None + try: + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + + query = "SELECT address, name FROM face WHERE id = %s" + cursor.execute(query, (face_id,)) + face = cursor.fetchone() + + if not face: + raise HTTPException(status_code=404, detail=f"ID为 {face_id} 的记录不存在") + + db_path = face["address"] + abs_path = Path(db_path).resolve() # 转为绝对路径 + if not db_path or not abs_path.exists(): + raise HTTPException(status_code=404, detail=f"图片不存在（路径：{db_path}）") + + return FileResponse( + path=abs_path, + filename=f"face_{face_id}_{face['name'] or '未命名'}.{db_path.split('.')[-1]}", + media_type=f"image/{db_path.split('.')[-1]}" + ) + except MySQLError as e: + raise HTTPException(status_code=500, detail=f"获取图片失败: {str(e)}") from e + finally: + db.close_connection(conn, cursor) + + +# ------------------------------ +# 内部工具方法 +# ------------------------------ +def get_all_face_name_with_eigenvalue() -> dict: conn = None cursor = None try: - # 1. 建立数据库连接并获取游标（dictionary=True使结果以字典形式返回） conn = db.get_connection() cursor = conn.cursor(dictionary=True) - # 2. 执行SQL查询: 只获取name非空的记录、减少数据传输 query = "SELECT name, eigenvalue FROM face WHERE name IS NOT NULL" cursor.execute(query) - faces = cursor.fetchall() # 返回结果: 列表套字典、如 [{"name":"张三","eigenvalue":...}, ...] + faces = cursor.fetchall() - # 3. 收集同一名称对应的所有特征值（处理名称重复场景） name_to_eigenvalues = {} for face in faces: name = face["name"] eigenvalue = face["eigenvalue"] - # 若名称已存在、追加特征值；否则新建列表存储 if name in name_to_eigenvalues: name_to_eigenvalues[name].append(eigenvalue) else: name_to_eigenvalues[name] = [eigenvalue] - # 4. 构建最终字典: 重复名称取平均、唯一名称直接取特征值 face_dict = {} for name, eigenvalues in name_to_eigenvalues.items(): - - # 处理特征值: 多个则求平均、单个则直接使用 if len(eigenvalues) > 1: - # 调用外部方法计算平均特征值（需确保binary_face_feature_handler已正确导入） face_dict[name] = get_average_feature(eigenvalues) else: - # 取列表中唯一的特征值（避免value为列表类型） face_dict[name] = eigenvalues[0] return face_dict except MySQLError as e: - # 捕获数据库异常、添加上下文信息后重新抛出（便于定位问题） - raise Exception(f"获取人脸名称与特征值失败: {str(e)}") from e + raise Exception(f"获取人脸特征失败: {str(e)}") from e finally: - # 5. 无论是否异常、均释放数据库连接和游标（避免资源泄漏） db.close_connection(conn, cursor) \ No newline at end of file diff --git a/service/model_service.py b/service/model_service.py new file mode 100644 index 0000000..8921c4a --- /dev/null +++ b/service/model_service.py @@ -0,0 +1,497 @@ +from fastapi import APIRouter, HTTPException, UploadFile, File, Form, Query +from fastapi.responses import FileResponse +from mysql.connector import Error as MySQLError +import os +import shutil +from pathlib import Path +from datetime import datetime + +# 复用项目依赖 +from ds.db import db +from schema.model_schema import ( + ModelCreateRequest, + ModelUpdateRequest, + ModelResponse, + ModelListResponse +) +from schema.response_schema import APIResponse +from util.model_util import load_yolo_model # 使用修复后的模型加载工具 + +# 路径配置 +CURRENT_FILE_PATH = Path(__file__).resolve() +PROJECT_ROOT = CURRENT_FILE_PATH.parent.parent +MODEL_SAVE_ROOT = PROJECT_ROOT / "resource" / "models" +MODEL_SAVE_ROOT.mkdir(exist_ok=True, parents=True) +DB_PATH_PREFIX_TO_REMOVE = str(PROJECT_ROOT) + os.sep + +# 模型限制 +ALLOWED_MODEL_EXT = {"pt"} +MAX_MODEL_SIZE = 100 * 1024 * 1024 # 100MB + +# 全局模型变量 +global _yolo_model +_yolo_model = None + +router = APIRouter(prefix="/models", tags=["模型管理"]) + + +# 工具函数：验证模型路径 +def get_valid_model_abs_path(relative_path: str) -> str: + try: + relative_path = relative_path.replace("/", os.sep) + model_abs_path = PROJECT_ROOT / relative_path + model_abs_path = model_abs_path.resolve() + model_abs_path_str = str(model_abs_path) + + if not model_abs_path_str.startswith(str(MODEL_SAVE_ROOT)): + raise HTTPException( + status_code=400, + detail=f"模型路径非法！允许目录：{str(MODEL_SAVE_ROOT)}，当前路径：{model_abs_path_str}" + ) + + if not model_abs_path.exists(): + raise HTTPException( + status_code=404, + detail=f"模型文件不存在！路径：{model_abs_path_str}" + ) + + if not model_abs_path.is_file(): + raise HTTPException( + status_code=400, + detail=f"路径不是文件！路径：{model_abs_path_str}" + ) + + file_size = model_abs_path.stat().st_size + if file_size > MAX_MODEL_SIZE: + raise HTTPException( + status_code=400, + detail=f"模型文件过大（{file_size // 1024 // 1024}MB），超过限制{MAX_MODEL_SIZE // 1024 // 1024}MB" + ) + + file_ext = model_abs_path.suffix.lower() + if file_ext not in [f".{ext}" for ext in ALLOWED_MODEL_EXT]: + raise HTTPException( + status_code=400, + detail=f"模型格式非法（仅支持{ALLOWED_MODEL_EXT}）！当前格式：{file_ext}" + ) + + print(f"[模型路径校验] 成功！路径：{model_abs_path_str}，大小：{file_size // 1024}KB") + return model_abs_path_str + + except HTTPException as e: + raise e + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"路径处理失败：{str(e)}" + ) from e + + +# 1. 上传模型 +@router.post("", response_model=APIResponse, summary="上传YOLO模型（.pt格式）") +async def upload_model( + name: str = Form(..., description="模型名称"), + description: str = Form(None, description="模型描述"), + is_default: bool = Form(False, description="是否设为默认模型"), + file: UploadFile = File(..., description=f"YOLO模型文件（.pt，最大{MAX_MODEL_SIZE // 1024 // 1024}MB）") +): + conn = None + cursor = None + saved_file_path = None + try: + # 校验文件 + file_ext = file.filename.split(".")[-1].lower() if "." in file.filename else "" + if file_ext not in ALLOWED_MODEL_EXT: + raise HTTPException( + status_code=400, + detail=f"仅支持{ALLOWED_MODEL_EXT}格式，当前：{file_ext}" + ) + if file.size > MAX_MODEL_SIZE: + raise HTTPException( + status_code=400, + detail=f"文件过大！最大{MAX_MODEL_SIZE // 1024 // 1024}MB，当前{file.size // 1024 // 1024}MB" + ) + + # 保存文件 + timestamp = datetime.now().strftime("%Y%m%d%H%M%S") + safe_filename = f"model_{timestamp}_{file.filename.replace(' ', '_')}" + saved_file_path = MODEL_SAVE_ROOT / safe_filename + with open(saved_file_path, "wb") as f: + shutil.copyfileobj(file.file, f) + saved_file_path.chmod(0o644) # 设置权限 + + # 数据库路径处理 + db_relative_path = str(saved_file_path).replace(DB_PATH_PREFIX_TO_REMOVE, "").replace(os.sep, "/") + + # 数据库操作 + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + + if is_default: + cursor.execute("UPDATE model SET is_default = 0") + + insert_sql = """ + INSERT INTO model (name, path, is_default, description, file_size) + VALUES (%s, %s, %s, %s, %s) + """ + cursor.execute(insert_sql, (name, db_relative_path, 1 if is_default else 0, description, file.size)) + conn.commit() + + cursor.execute("SELECT * FROM model WHERE id = LAST_INSERT_ID()") + new_model = cursor.fetchone() + if not new_model: + raise HTTPException(status_code=500, detail="上传成功但无法获取记录") + + # 加载默认模型 + global _yolo_model + if is_default: + valid_abs_path = get_valid_model_abs_path(db_relative_path) + _yolo_model = load_yolo_model(valid_abs_path) + if not _yolo_model: + raise HTTPException( + status_code=500, + detail=f"上传成功，但加载默认模型失败（路径：{valid_abs_path}）" + ) + + return APIResponse( + code=201, + message=f"模型上传成功！ID：{new_model['id']}", + data=ModelResponse(**new_model) + ) + + except MySQLError as e: + if conn: + conn.rollback() + if saved_file_path and saved_file_path.exists(): + saved_file_path.unlink() + raise HTTPException(status_code=500, detail=f"数据库错误：{str(e)}") from e + except Exception as e: + if saved_file_path and saved_file_path.exists(): + saved_file_path.unlink() + raise HTTPException(status_code=500, detail=f"服务器错误：{str(e)}") from e + finally: + await file.close() + db.close_connection(conn, cursor) + + +# 2. 获取模型列表 +@router.get("", response_model=APIResponse, summary="获取模型列表（分页）") +async def get_model_list( + page: int = Query(1, ge=1), + page_size: int = Query(10, ge=1, le=100), + name: str = Query(None), + is_default: bool = Query(None) +): + conn = None + cursor = None + try: + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + + where_clause = [] + params = [] + if name: + where_clause.append("name LIKE %s") + params.append(f"%{name}%") + if is_default is not None: + where_clause.append("is_default = %s") + params.append(1 if is_default else 0) + + # 总记录数 + count_sql = "SELECT COUNT(*) AS total FROM model" + if where_clause: + count_sql += " WHERE " + " AND ".join(where_clause) + cursor.execute(count_sql, params) + total = cursor.fetchone()["total"] + + # 分页数据 + offset = (page - 1) * page_size + list_sql = "SELECT * FROM model" + if where_clause: + list_sql += " WHERE " + " AND ".join(where_clause) + list_sql += " ORDER BY updated_at DESC LIMIT %s OFFSET %s" + params.extend([page_size, offset]) + + cursor.execute(list_sql, params) + model_list = cursor.fetchall() + + return APIResponse( + code=200, + message=f"获取成功！共{total}条记录", + data=ModelListResponse( + total=total, + models=[ModelResponse(**model) for model in model_list] + ) + ) + + except MySQLError as e: + raise HTTPException(status_code=500, detail=f"数据库错误：{str(e)}") from e + finally: + db.close_connection(conn, cursor) + + +# 3. 获取默认模型 +@router.get("/default", response_model=APIResponse, summary="获取当前默认模型") +async def get_default_model(): + conn = None + cursor = None + try: + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + + cursor.execute("SELECT * FROM model WHERE is_default = 1") + default_model = cursor.fetchone() + + if not default_model: + raise HTTPException(status_code=404, detail="暂无默认模型") + + valid_abs_path = get_valid_model_abs_path(default_model["path"]) + global _yolo_model + + if not _yolo_model: + _yolo_model = load_yolo_model(valid_abs_path) + if not _yolo_model: + raise HTTPException( + status_code=500, + detail=f"默认模型存在，但加载失败（路径：{valid_abs_path}）" + ) + + return APIResponse( + code=200, + message="默认模型查询成功", + data=ModelResponse(**default_model) + ) + + except MySQLError as e: + raise HTTPException(status_code=500, detail=f"数据库错误：{str(e)}") from e + finally: + db.close_connection(conn, cursor) + + +# 4. 获取单个模型详情 +@router.get("/{model_id}", response_model=APIResponse, summary="获取单个模型详情") +async def get_model(model_id: int): + conn = None + cursor = None + try: + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + + cursor.execute("SELECT * FROM model WHERE id = %s", (model_id,)) + model = cursor.fetchone() + + if not model: + raise HTTPException(status_code=404, detail=f"模型不存在！ID：{model_id}") + + try: + model_abs_path = get_valid_model_abs_path(model["path"]) + except HTTPException as e: + return APIResponse( + code=200, + message=f"查询成功，但路径异常：{e.detail}", + data=ModelResponse(**model) + ) + + return APIResponse( + code=200, + message="查询成功", + data=ModelResponse(**model) + ) + + except MySQLError as e: + raise HTTPException(status_code=500, detail=f"数据库错误：{str(e)}") from e + finally: + db.close_connection(conn, cursor) + + +# 5. 更新模型信息 +@router.put("/{model_id}", response_model=APIResponse, summary="更新模型信息") +async def update_model(model_id: int, model_update: ModelUpdateRequest): + conn = None + cursor = None + try: + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + + cursor.execute("SELECT * FROM model WHERE id = %s", (model_id,)) + exist_model = cursor.fetchone() + if not exist_model: + raise HTTPException(status_code=404, detail=f"模型不存在！ID：{model_id}") + + update_fields = [] + params = [] + if model_update.name is not None: + update_fields.append("name = %s") + params.append(model_update.name) + if model_update.description is not None: + update_fields.append("description = %s") + params.append(model_update.description) + + need_load_default = False + if model_update.is_default is not None: + if model_update.is_default: + cursor.execute("UPDATE model SET is_default = 0") + update_fields.append("is_default = 1") + need_load_default = True + else: + cursor.execute("SELECT COUNT(*) AS cnt FROM model WHERE is_default = 1") + default_count = cursor.fetchone()["cnt"] + if default_count == 1 and exist_model["is_default"]: + raise HTTPException( + status_code=400, + detail="当前是唯一默认模型，不可取消！" + ) + update_fields.append("is_default = 0") + + if not update_fields: + raise HTTPException(status_code=400, detail="至少需提供一个更新字段") + + params.append(model_id) + update_sql = f""" + UPDATE model + SET {', '.join(update_fields)}, updated_at = CURRENT_TIMESTAMP + WHERE id = %s + """ + cursor.execute(update_sql, params) + conn.commit() + + cursor.execute("SELECT * FROM model WHERE id = %s", (model_id,)) + updated_model = cursor.fetchone() + + global _yolo_model + if need_load_default: + valid_abs_path = get_valid_model_abs_path(updated_model["path"]) + _yolo_model = load_yolo_model(valid_abs_path) + if not _yolo_model: + raise HTTPException( + status_code=500, + detail=f"更新成功，但加载新默认模型失败（路径：{valid_abs_path}）" + ) + + return APIResponse( + code=200, + message="模型更新成功", + data=ModelResponse(**updated_model) + ) + + except MySQLError as e: + if conn: + conn.rollback() + raise HTTPException(status_code=500, detail=f"数据库错误：{str(e)}") from e + finally: + db.close_connection(conn, cursor) + + +# 6. 删除模型 +@router.delete("/{model_id}", response_model=APIResponse, summary="删除模型") +async def delete_model(model_id: int): + conn = None + cursor = None + try: + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + + cursor.execute("SELECT * FROM model WHERE id = %s", (model_id,)) + exist_model = cursor.fetchone() + if not exist_model: + raise HTTPException(status_code=404, detail=f"模型不存在！ID：{model_id}") + if exist_model["is_default"]: + raise HTTPException(status_code=400, detail="默认模型不可删除！") + + try: + model_abs_path_str = get_valid_model_abs_path(exist_model["path"]) + model_abs_path = Path(model_abs_path_str) + except HTTPException as e: + cursor.execute("DELETE FROM model WHERE id = %s", (model_id,)) + conn.commit() + return APIResponse( + code=200, + message=f"记录删除成功，文件异常：{e.detail}", + data=None + ) + + cursor.execute("DELETE FROM model WHERE id = %s", (model_id,)) + conn.commit() + + extra_msg = "" + try: + model_abs_path.unlink() + extra_msg = f"（已删除文件）" + except Exception as e: + extra_msg = f"（文件删除失败：{str(e)}）" + + global _yolo_model + if _yolo_model and str(_yolo_model.model_path) == model_abs_path_str: + _yolo_model = None + print(f"[模型删除] 已清空全局模型（路径：{model_abs_path_str}）") + + return APIResponse( + code=200, + message=f"模型删除成功！ID：{model_id} {extra_msg}", + data=None + ) + + except MySQLError as e: + if conn: + conn.rollback() + raise HTTPException(status_code=500, detail=f"数据库错误：{str(e)}") from e + finally: + db.close_connection(conn, cursor) + + +# 7. 下载模型文件 +@router.get("/{model_id}/download", summary="下载模型文件") +async def download_model(model_id: int): + conn = None + cursor = None + try: + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + + cursor.execute("SELECT * FROM model WHERE id = %s", (model_id,)) + model = cursor.fetchone() + if not model: + raise HTTPException(status_code=404, detail=f"模型不存在！ID：{model_id}") + + valid_abs_path = get_valid_model_abs_path(model["path"]) + model_abs_path = Path(valid_abs_path) + + return FileResponse( + path=model_abs_path, + filename=f"model_{model_id}_{model['name']}.pt", + media_type="application/octet-stream" + ) + + except MySQLError as e: + raise HTTPException(status_code=500, detail=f"数据库错误：{str(e)}") from e + finally: + db.close_connection(conn, cursor) + + +# 对外提供当前模型 +def get_current_yolo_model(): + """供检测模块获取当前加载的模型""" + global _yolo_model + if not _yolo_model: + conn = None + cursor = None + try: + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + cursor.execute("SELECT path FROM model WHERE is_default = 1") + default_model = cursor.fetchone() + if not default_model: + print("[get_current_yolo_model] 暂无默认模型") + return None + + valid_abs_path = get_valid_model_abs_path(default_model["path"]) + _yolo_model = load_yolo_model(valid_abs_path) + if _yolo_model: + print(f"[get_current_yolo_model] 自动加载默认模型成功") + else: + print(f"[get_current_yolo_model] 自动加载默认模型失败") + except Exception as e: + print(f"[get_current_yolo_model] 加载失败：{str(e)}") + finally: + db.close_connection(conn, cursor) + return _yolo_model diff --git a/service/user_service.py b/service/user_service.py index f3c5da7..f96dcdb 100644 --- a/service/user_service.py +++ b/service/user_service.py @@ -1,6 +1,7 @@ from datetime import timedelta +from typing import Optional -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException, Query from mysql.connector import Error as MySQLError from ds.db import db @@ -11,7 +12,7 @@ from middle.auth_middleware import ( verify_password, create_access_token, ACCESS_TOKEN_EXPIRE_MINUTES, - get_current_user + get_current_user # 仅保留登录用户校验，移除is_admin导入 ) # 创建用户接口路由（前缀 /users、标签用于 Swagger 分类） @@ -27,7 +28,7 @@ router = APIRouter( @router.post("/register", response_model=APIResponse, summary="用户注册") async def user_register(request: UserRegisterRequest): """ - 用户注册: + 用户注册: - 校验用户名是否已存在 - 加密密码后插入数据库 - 返回注册成功信息 @@ -78,7 +79,7 @@ async def user_register(request: UserRegisterRequest): @router.post("/login", response_model=APIResponse, summary="用户登录（获取 Token）") async def user_login(request: UserLoginRequest): """ - 用户登录: + 用户登录: - 校验用户名是否存在 - 校验密码是否正确 - 生成 JWT Token 并返回 @@ -142,7 +143,7 @@ async def get_current_user_info( current_user: UserResponse = Depends(get_current_user) # 依赖认证中间件 ): """ - 获取当前登录用户信息: + 获取当前登录用户信息: - 需在请求头携带 Token（格式: Bearer ） - 认证通过后返回用户信息 """ @@ -152,3 +153,98 @@ async def get_current_user_info( data=current_user ) + +# ------------------------------ +# 4. 获取用户列表（仅需登录权限） +# ------------------------------ +@router.get("/list", response_model=APIResponse, summary="获取用户列表") +async def get_user_list( + page: int = Query(1, ge=1, description="页码，从1开始"), + page_size: int = Query(10, ge=1, le=100, description="每页条数，1-100之间"), + username: Optional[str] = Query(None, description="用户名模糊搜索"), + current_user: UserResponse = Depends(get_current_user) # 仅需登录即可访问（移除管理员校验） +): + """ + 获取用户列表: + - 需登录权限（请求头携带 Token: Bearer ） + - 支持分页查询（page=页码，page_size=每页条数） + - 支持用户名模糊搜索（如输入"test"可匹配"test123"、"admin_test"等） + - 仅返回用户ID、用户名、创建时间、更新时间（不包含密码等敏感信息） + """ + conn = None + cursor = None + try: + conn = db.get_connection() + cursor = conn.cursor(dictionary=True) + + # 计算分页偏移量（page从1开始，偏移量=（页码-1）*每页条数） + offset = (page - 1) * page_size + + # 基础查询（仅查非敏感字段） + base_query = """ + SELECT id, username, created_at, updated_at + FROM users + """ + # 总条数查询（用于分页计算） + count_query = "SELECT COUNT(*) as total FROM users" + + # 条件拼接（支持用户名模糊搜索） + conditions = [] + params = [] + if username: + conditions.append("username LIKE %s") + params.append(f"%{username}%") # 模糊匹配：%表示任意字符 + + # 构建最终查询语句 + if conditions: + where_clause = " WHERE " + " AND ".join(conditions) + final_query = f"{base_query}{where_clause} LIMIT %s OFFSET %s" + final_count_query = f"{count_query}{where_clause}" + params.extend([page_size, offset]) # 追加分页参数 + else: + final_query = f"{base_query} LIMIT %s OFFSET %s" + final_count_query = count_query + params = [page_size, offset] + + # 1. 查询用户列表数据 + cursor.execute(final_query, params) + users = cursor.fetchall() + + # 2. 查询总条数（用于计算总页数） + count_params = [f"%{username}%"] if username else [] + cursor.execute(final_count_query, count_params) + total = cursor.fetchone()["total"] + + # 3. 转换为UserResponse模型（确保字段匹配） + user_list = [ + UserResponse( + id=user["id"], + username=user["username"], + created_at=user["created_at"], + updated_at=user["updated_at"] + ) + for user in users + ] + + # 4. 计算总页数（向上取整，如11条数据每页10条=2页） + total_pages = (total + page_size - 1) // page_size + + # 返回结果（包含列表和分页信息） + return APIResponse( + code=200, + message="获取用户列表成功", + data={ + "users": user_list, + "pagination": { + "page": page, # 当前页码 + "page_size": page_size, # 每页条数 + "total": total, # 总数据量 + "total_pages": total_pages # 总页数 + } + } + ) + 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/util/face_util.py b/util/face_util.py index 54bf989..2e0470f 100644 --- a/util/face_util.py +++ b/util/face_util.py @@ -4,6 +4,11 @@ import insightface from insightface.app import FaceAnalysis from io import BytesIO from PIL import Image +import logging + +# 配置日志（便于排查） +logging.basicConfig(level=logging.INFO, format='%(asctime)s - [FaceUtil] - %(levelname)s - %(message)s') +logger = logging.getLogger(__name__) # 全局变量存储InsightFace引擎和特征列表 _insightface_app = None @@ -11,135 +16,141 @@ _feature_list = [] def init_insightface(): - """初始化InsightFace引擎""" + """初始化InsightFace引擎（确保成功后再使用）""" global _insightface_app try: - print("正在初始化InsightFace引擎...") - app = FaceAnalysis(name='buffalo_l', root='~/.insightface') - app.prepare(ctx_id=0, det_size=(640, 640)) - print("InsightFace引擎初始化完成") + if _insightface_app is not None: + logger.info("InsightFace引擎已初始化，无需重复执行") + return _insightface_app + + logger.info("正在初始化InsightFace引擎（模型：buffalo_l）...") + # 手动指定模型下载路径（避免权限问题，可选） + app = FaceAnalysis( + name='buffalo_l', + root='~/.insightface', # 模型默认下载路径 + providers=['CPUExecutionProvider'] # 强制用CPU（若有GPU可加'CUDAExecutionProvider'） + ) + app.prepare(ctx_id=0, det_size=(640, 640)) # det_size越大，小人脸检测越准 + logger.info("InsightFace引擎初始化完成") _insightface_app = app return app except Exception as e: - print(f"InsightFace初始化失败: {e}") + logger.error(f"InsightFace初始化失败：{str(e)}", exc_info=True) # 打印详细堆栈 + _insightface_app = None return None def add_binary_data(binary_data): """ 接收单张图片的二进制数据、提取特征并保存 - - 参数: - binary_data: 图片的二进制数据（bytes类型） - - 返回: - 成功提取特征时返回 (True, 特征值numpy数组) - 失败时返回 (False, None) + 返回：(True, 特征值numpy数组) 或 (False, 错误信息字符串) """ global _insightface_app, _feature_list + # 1. 先检查引擎是否初始化成功 if not _insightface_app: - print("引擎未初始化、无法处理") - return False, None + init_result = init_insightface() # 尝试重新初始化 + if not init_result: + error_msg = "InsightFace引擎未初始化，无法检测人脸" + logger.error(error_msg) + return False, error_msg try: - # 直接处理二进制数据: 转换为图像格式 - img = Image.open(BytesIO(binary_data)) - frame = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) + # 2. 验证二进制数据有效性 + if len(binary_data) < 1024: # 过滤过小的无效图片（小于1KB） + error_msg = f"图片过小（{len(binary_data)}字节），可能不是有效图片" + logger.warning(error_msg) + return False, error_msg - # 提取特征 + # 3. 二进制数据转CV2格式（关键步骤，避免通道错误） + try: + img = Image.open(BytesIO(binary_data)).convert("RGB") # 强制转RGB + frame = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) # InsightFace需要BGR格式 + except Exception as e: + error_msg = f"图片格式转换失败：{str(e)}" + logger.error(error_msg, exc_info=True) + return False, error_msg + + # 4. 检查图片尺寸（避免极端尺寸导致检测失败） + height, width = frame.shape[:2] + if height < 64 or width < 64: # 人脸检测最小建议尺寸 + error_msg = f"图片尺寸过小（{width}x{height}），需至少64x64像素" + logger.warning(error_msg) + return False, error_msg + + # 5. 调用InsightFace检测人脸 + logger.info(f"开始检测人脸（图片尺寸：{width}x{height}，格式：BGR）") faces = _insightface_app.get(frame) - if faces: - # 获取当前提取的特征值 - current_feature = faces[0].embedding - # 添加到特征列表 - _feature_list.append(current_feature) - print(f"已累计 {len(_feature_list)} 个特征") - # 返回成功标志和当前特征值 - return True, current_feature - else: - print("二进制数据中未检测到人脸") - return False, None + + if not faces: + error_msg = "未检测到人脸（请确保图片包含清晰正面人脸，无遮挡、不模糊）" + logger.warning(error_msg) + return False, error_msg + + # 6. 提取特征并保存 + current_feature = faces[0].embedding + _feature_list.append(current_feature) + logger.info(f"人脸检测成功，提取特征值（维度：{current_feature.shape[0]}），累计特征数：{len(_feature_list)}") + return True, current_feature + except Exception as e: - print(f"处理二进制数据出错: {e}") - return False, None + error_msg = f"处理图片时发生异常：{str(e)}" + logger.error(error_msg, exc_info=True) + return False, error_msg +# 以下函数保持不变（get_average_feature/clear_features/get_feature_list） def get_average_feature(features=None): - """ - 计算多个特征向量的平均值 - - 参数: - features: 可选、特征值列表。如果未提供、则使用全局存储的_feature_list - 每个元素可以是字符串格式或numpy数组 - - 返回: - 单一平均特征向量的numpy数组、若无可计算数据则返回None - """ global _feature_list - - # 如果未提供features参数、则使用全局特征列表 - if features is None: - features = _feature_list - try: - # 验证输入是否为列表且不为空 + if features is None: + features = _feature_list if not isinstance(features, list) or len(features) == 0: - print("输入必须是包含至少一个特征值的列表") + logger.warning("输入必须是包含至少一个特征值的列表") return None - # 处理每个特征值 processed_features = [] for i, embedding in enumerate(features): try: if isinstance(embedding, str): - # 处理包含括号和逗号的字符串格式 embedding_str = embedding.replace('[', '').replace(']', '').replace(',', ' ').strip() embedding_list = [float(num) for num in embedding_str.split() if num.strip()] embedding_np = np.array(embedding_list, dtype=np.float32) else: embedding_np = np.array(embedding, dtype=np.float32) - # 验证特征值格式 if len(embedding_np.shape) == 1: processed_features.append(embedding_np) - print(f"已添加第 {i + 1} 个特征值用于计算平均值") + logger.info(f"已添加第 {i + 1} 个特征值用于计算平均值") else: - print(f"跳过第 {i + 1} 个特征值、不是一维数组") - + logger.warning(f"跳过第 {i + 1} 个特征值：不是一维数组") except Exception as e: - print(f"处理第 {i + 1} 个特征值时出错: {e}") + logger.error(f"处理第 {i + 1} 个特征值时出错：{str(e)}") - # 确保有有效的特征值 if not processed_features: - print("没有有效的特征值用于计算平均值") + logger.warning("没有有效的特征值用于计算平均值") return None - # 检查所有特征向量维度是否相同 dims = {feat.shape[0] for feat in processed_features} if len(dims) > 1: - print(f"特征值维度不一致、无法计算平均值。检测到的维度: {dims}") + logger.error(f"特征值维度不一致：{dims}，无法计算平均值") return None - # 计算平均值 avg_feature = np.mean(processed_features, axis=0) - print(f"成功计算 {len(processed_features)} 个特征值的平均特征向量、维度: {avg_feature.shape[0]}") - + logger.info(f"计算成功：{len(processed_features)} 个特征值的平均向量（维度：{avg_feature.shape[0]}）") return avg_feature - except Exception as e: - print(f"计算平均特征值时出错: {e}") + logger.error(f"计算平均特征值出错：{str(e)}", exc_info=True) return None def clear_features(): - """清空已存储的特征数据""" global _feature_list _feature_list = [] - print("已清空所有特征数据") + logger.info("已清空所有特征数据") def get_feature_list(): - """获取当前存储的特征列表""" global _feature_list - return _feature_list.copy() # 返回副本防止外部直接修改 \ No newline at end of file + logger.info(f"当前特征列表长度：{len(_feature_list)}") + return _feature_list.copy() \ No newline at end of file diff --git a/util/file_util.py b/util/file_util.py new file mode 100644 index 0000000..8e60e9d --- /dev/null +++ b/util/file_util.py @@ -0,0 +1,83 @@ +import os +import datetime +from pathlib import Path +from typing import Dict + + +def save_face_to_up_images( + client_ip: str, + face_name: str, + image_bytes: bytes, + image_format: str = "jpg" +) -> Dict[str, str]: + """ + 保存人脸图片到 `/up_images/用户IP/人脸名字/` 路径 + 修复路径计算错误，确保所有路径在up_images根目录下 + + 参数: + client_ip: 客户端IP（原始格式，如192.168.1.101） + face_name: 人脸名称（用户输入，可为空） + image_bytes: 人脸图片二进制数据 + image_format: 图片格式（默认jpg） + + 返回: + 字典：success（是否成功）、db_path（存数据库的相对路径）、local_abs_path（本地绝对路径）、msg（提示） + """ + try: + # 1. 基础参数校验 + if not client_ip.strip(): + return {"success": False, "db_path": "", "local_abs_path": "", "msg": "客户端IP不能为空"} + if not image_bytes: + return {"success": False, "db_path": "", "local_abs_path": "", "msg": "图片二进制数据为空"} + if image_format.lower() not in ["jpg", "jpeg", "png"]: + return {"success": False, "db_path": "", "local_abs_path": "", "msg": "仅支持jpg/jpeg/png格式"} + + # 2. 处理特殊字符（避免路径错误） + safe_ip = client_ip.strip().replace(".", "_") # IP中的.替换为_ + safe_face_name = face_name.strip() if (face_name and face_name.strip()) else "未命名" + safe_face_name = "".join([c for c in safe_face_name if c not in r'\/:*?"<>|']) # 过滤非法字符 + + # 3. 构建根目录（强制转为绝对路径，避免相对路径混淆） + root_dir = Path("up_images").resolve() # 转为绝对路径（关键修复！） + if not root_dir.exists(): + root_dir.mkdir(parents=True, exist_ok=True) + print(f"[FileUtil] 已创建up_images根目录：{root_dir}") + + # 4. 构建文件层级路径（确保在root_dir子目录下） + ip_dir = root_dir / safe_ip + face_name_dir = ip_dir / safe_face_name + face_name_dir.mkdir(parents=True, exist_ok=True) # 自动创建目录 + print(f"[FileUtil] 图片存储目录：{face_name_dir}") + + # 5. 生成唯一文件名（毫秒级时间戳） + timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S%f")[:-3] + image_filename = f"face_{safe_ip}_{timestamp}.{image_format.lower()}" + + # 6. 计算路径（关键修复：确保所有路径都是绝对路径且在root_dir下） + local_abs_path = face_name_dir / image_filename # 绝对路径 + + # 验证路径是否在root_dir下（防止路径穿越攻击） + if not local_abs_path.resolve().is_relative_to(root_dir.resolve()): + raise Exception(f"图片路径不在up_images根目录下（安全校验失败）：{local_abs_path}") + + # 数据库存储路径：从root_dir开始的相对路径（如 up_images/192_168_110_31/小王/xxx.jpg） + db_path = str(root_dir.name / local_abs_path.relative_to(root_dir)) + + # 7. 写入图片文件 + with open(local_abs_path, "wb") as f: + f.write(image_bytes) + print(f"[FileUtil] 图片保存成功：") + print(f" 数据库路径：{db_path}") + print(f" 本地绝对路径：{local_abs_path}") + + return { + "success": True, + "db_path": db_path, # 存数据库的相对路径（up_images开头） + "local_abs_path": str(local_abs_path), # 本地绝对路径 + "msg": "图片保存成功" + } + + except Exception as e: + error_msg = f"图片保存失败：{str(e)}" + print(f"[FileUtil] 错误：{error_msg}") + return {"success": False, "db_path": "", "local_abs_path": "", "msg": error_msg} \ No newline at end of file diff --git a/util/model_util.py b/util/model_util.py new file mode 100644 index 0000000..1b365ea --- /dev/null +++ b/util/model_util.py @@ -0,0 +1,61 @@ +import os +import numpy as np +import traceback +from ultralytics import YOLO +from typing import Optional + + +def load_yolo_model(model_path: str) -> Optional[YOLO]: + """ + 加载YOLO模型（支持v5/v8），并校验模型有效性 + :param model_path: 模型文件的绝对路径 + :return: 加载成功返回YOLO模型实例，失败返回None + """ + try: + # 加载前的基础信息检查 + print(f"\n[模型工具] 开始加载模型：{model_path}") + print(f"[模型工具] 文件是否存在：{os.path.exists(model_path)}") + if os.path.exists(model_path): + print(f"[模型工具] 文件大小：{os.path.getsize(model_path) / 1024 / 1024:.2f} MB") + + # 强制重新加载模型，避免缓存问题 + model = YOLO(model_path) + + # 兼容性校验：使用numpy空数组测试模型 + dummy_image = np.zeros((640, 640, 3), dtype=np.uint8) + + try: + # 优先使用新版本参数 + model.predict( + source=dummy_image, + imgsz=640, + conf=0.25, + verbose=False, + stream=False + ) + except Exception as pred_e: + print(f"[模型工具] 预测校验兼容处理：{str(pred_e)}") + # 兼容旧版本YOLO参数 + model.predict( + img=dummy_image, + imgsz=640, + conf=0.25, + verbose=False + ) + + # 验证模型基本属性 + if not hasattr(model, 'names'): + print("[模型工具] 警告：模型缺少类别名称属性") + else: + print(f"[模型工具] 模型包含类别：{list(model.names.values())[:5]}...") # 显示前5个类别 + + print(f"[模型工具] 模型加载成功！") + return model + + except Exception as e: + # 详细错误信息输出 + print(f"\n[模型工具] 加载模型失败！路径：{model_path}") + print(f"[模型工具] 异常类型：{type(e).__name__}") + print(f"[模型工具] 异常详情：{str(e)}") + print(f"[模型工具] 堆栈跟踪：\n{traceback.format_exc()}") + return None