283 lines
12 KiB
Python
283 lines
12 KiB
Python
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/<path:resource_path>')
|
||
@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/<path:resource_path>')
|
||
@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/<path:resource_path>')
|
||
@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/<path:resource_path>')
|
||
@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
|
||
) |