Spaces:

ethonmax
/

picpocket

Running

chawin.chen

web

7396c20 11 months ago

7.57 kB

import os
import time
from contextlib import asynccontextmanager

import base64
import secrets

from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from starlette.middleware.cors import CORSMiddleware
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response, PlainTextResponse

from cleanup_scheduler import start_cleanup_scheduler, stop_cleanup_scheduler
from config import (
    logger,
    OUTPUT_DIR,
    DEEPFACE_AVAILABLE,
    DLIB_AVAILABLE,
    MODELS_PATH,
    IMAGES_DIR,
    YOLO_AVAILABLE,
    ENABLE_LOGGING,
    HUGGINGFACE_SYNC_ENABLED,
)
from database import close_mysql_pool, init_mysql_pool
from utils import ensure_bos_resources, ensure_huggingface_models

logger.info("Starting to import api_routes module...")

if HUGGINGFACE_SYNC_ENABLED:
    try:
        t_hf_start = time.perf_counter()
        if not ensure_huggingface_models():
            raise RuntimeError("无法从 HuggingFace 同步模型，请检查配置与网络")
        hf_time = time.perf_counter() - t_hf_start
        logger.info("HuggingFace 模型同步完成，用时 %.3fs", hf_time)
    except Exception as exc:
        logger.error(f"HuggingFace model preparation failed: {exc}")
        raise
else:
    logger.info("已关闭 HuggingFace 模型同步开关，跳过启动阶段的同步步骤")

try:
    t_bos_start = time.perf_counter()
    if not ensure_bos_resources():
        raise RuntimeError("无法从 BOS 同步模型与数据，请检查凭证与网络")
    bos_time = time.perf_counter() - t_bos_start
    logger.info(f"BOS resources synchronized successfully, time: {bos_time:.3f}s")
except Exception as exc:
    logger.error(f"BOS resource preparation failed: {exc}")
    raise

try:
    t_start = time.perf_counter()
    from api_routes import api_router, extract_chinese_celeb_dataset_sync
    import_time = time.perf_counter() - t_start
    logger.info(f"api_routes module imported successfully, time: {import_time:.3f}s")
except Exception as e:
    import_time = time.perf_counter() - t_start
    logger.error(f"api_routes module import failed, time: {import_time:.3f}s, error: {e}")
    raise

try:
    t_extract_start = time.perf_counter()
    extract_result = extract_chinese_celeb_dataset_sync()
    extract_time = time.perf_counter() - t_extract_start
    logger.info(
        "Chinese celeb dataset extracted successfully, time: %.3fs, target: %s",
        extract_time,
        extract_result.get("target_dir"),
    )
except Exception as exc:
    logger.error(f"Failed to extract Chinese celeb dataset automatically: {exc}")
    raise


@asynccontextmanager
async def lifespan(app: FastAPI):
    start_time = time.perf_counter()
    logger.info("FaceScore service starting...")
    logger.info(f"Output directory: {OUTPUT_DIR}")
    logger.info(f"DeepFace available: {DEEPFACE_AVAILABLE}")
    logger.info(f"YOLO available: {YOLO_AVAILABLE}")
    logger.info(f"MediaPipe available: {DLIB_AVAILABLE}")
    logger.info(f"Archive directory: {IMAGES_DIR}")
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # 初始化数据库连接池
    try:
      await init_mysql_pool()
      logger.info("MySQL 连接池初始化完成")
    except Exception as exc:
      logger.error(f"初始化 MySQL 连接池失败: {exc}")
      raise

    # 启动图片清理定时任务
    logger.info("Starting image cleanup scheduled task...")
    try:
        start_cleanup_scheduler()
        logger.info("Image cleanup scheduled task started successfully")
    except Exception as e:
        logger.error(f"Failed to start image cleanup scheduled task: {e}")

    # 记录启动完成时间
    total_startup_time = time.perf_counter() - start_time
    logger.info(f"FaceScore service startup completed, total time: {total_startup_time:.3f}s")

    yield

    # 应用关闭时停止定时任务
    logger.info("Stopping image cleanup scheduled task...")
    try:
        stop_cleanup_scheduler()
        logger.info("Image cleanup scheduled task stopped")
    except Exception as e:
        logger.error(f"Failed to stop image cleanup scheduled task: {e}")

    # 关闭数据库连接池
    try:
      await close_mysql_pool()
    except Exception as exc:
      logger.warning(f"关闭 MySQL 连接池失败: {exc}")


# 创建 FastAPI 应用
app = FastAPI(
    title="Enhanced FaceScore 服务",
    description="支持多模型的人脸分析REST API服务，包含五官评分功能。支持混合模式：HowCuteAmI（颜值+性别）+ DeepFace（年龄+情绪）",
    version="3.0.0",
    docs_url="/cp_docs",
    redoc_url="/cp_redoc",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# 注册路由
app.include_router(api_router)


_default_frontend_dir = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "facelist-web")
)
FRONTEND_DIR = os.path.abspath(
    os.path.expanduser(os.environ.get("FACELIST_FRONTEND_DIR", _default_frontend_dir))
)

_basic_user = os.environ.get("FACELIST_BASIC_USER", "admin")
_basic_pass = os.environ.get("FACELIST_BASIC_PASS", "admin")
_basic_secret = f"{_basic_user}:{_basic_pass}"
_basic_token = "Basic " + base64.b64encode(_basic_secret.encode()).decode()


class FacelistBasicAuthMiddleware(BaseHTTPMiddleware):
    """仅保护 /facelist 前缀的 Basic Auth"""

    async def dispatch(self, request, call_next):
        path = request.url.path
        if path.startswith("/facelist"):
            auth = request.headers.get("Authorization", "")
            if not secrets.compare_digest(auth, _basic_token):
                return PlainTextResponse(
                    "Unauthorized",
                    status_code=401,
                    headers={"WWW-Authenticate": 'Basic realm="facelist"'},
                )
        return await call_next(request)


class SPAStaticFiles(StaticFiles):
    """支持SPA路由回退到index.html的静态文件服务"""

    async def get_response(self, path: str, scope):
        response = await super().get_response(path, scope)
        if response.status_code == 404:
            # 未匹配到具体文件时回退到前端入口
            response = await super().get_response("index.html", scope)
        return response

app.add_middleware(FacelistBasicAuthMiddleware)

if os.path.isdir(FRONTEND_DIR):
    app.mount(
        "/facelist",
        SPAStaticFiles(directory=FRONTEND_DIR, html=True),
        name="facelist",
    )
    logger.info("前端静态资源已挂载在 /facelist ，目录: %s", FRONTEND_DIR)
else:
    logger.warning("未找到前端目录 %s ，跳过前端静态资源挂载", FRONTEND_DIR)

# 添加根路径处理
@app.get("/")
async def root():
    return "UP"


if __name__ == "__main__":
    import uvicorn

    if not os.path.exists(MODELS_PATH):
        logger.critical(
            "Warning: 'models' directory not found. Please ensure it exists and contains model files."
        )
        logger.critical(
            "Exiting application as FaceAnalyzer cannot be initialized without models."
        )
        exit(1)

    # 根据日志开关配置 Uvicorn 日志
    if ENABLE_LOGGING:
        uvicorn.run(app, host="0.0.0.0", port=8080, reload=False)
    else:
        # 禁用 Uvicorn 的访问日志和错误日志
        uvicorn.run(
            app,
            host="0.0.0.0",
            port=8080,
            reload=False,
            access_log=False,  # 禁用访问日志
            log_level="critical"  # 只显示严重错误
        )

	import os
	import time
	from contextlib import asynccontextmanager

	import base64
	import secrets

	from fastapi import FastAPI
	from fastapi.staticfiles import StaticFiles
	from starlette.middleware.cors import CORSMiddleware
	from starlette.middleware.base import BaseHTTPMiddleware
	from starlette.responses import Response, PlainTextResponse

	from cleanup_scheduler import start_cleanup_scheduler, stop_cleanup_scheduler
	from config import (
	logger,
	OUTPUT_DIR,
	DEEPFACE_AVAILABLE,
	DLIB_AVAILABLE,
	MODELS_PATH,
	IMAGES_DIR,
	YOLO_AVAILABLE,
	ENABLE_LOGGING,
	HUGGINGFACE_SYNC_ENABLED,
	)
	from database import close_mysql_pool, init_mysql_pool
	from utils import ensure_bos_resources, ensure_huggingface_models

	logger.info("Starting to import api_routes module...")

	if HUGGINGFACE_SYNC_ENABLED:
	try:
	t_hf_start = time.perf_counter()
	if not ensure_huggingface_models():
	raise RuntimeError("无法从 HuggingFace 同步模型，请检查配置与网络")
	hf_time = time.perf_counter() - t_hf_start
	logger.info("HuggingFace 模型同步完成，用时 %.3fs", hf_time)
	except Exception as exc:
	logger.error(f"HuggingFace model preparation failed: {exc}")
	raise
	else:
	logger.info("已关闭 HuggingFace 模型同步开关，跳过启动阶段的同步步骤")

	try:
	t_bos_start = time.perf_counter()
	if not ensure_bos_resources():
	raise RuntimeError("无法从 BOS 同步模型与数据，请检查凭证与网络")
	bos_time = time.perf_counter() - t_bos_start
	logger.info(f"BOS resources synchronized successfully, time: {bos_time:.3f}s")
	except Exception as exc:
	logger.error(f"BOS resource preparation failed: {exc}")
	raise

	try:
	t_start = time.perf_counter()
	from api_routes import api_router, extract_chinese_celeb_dataset_sync
	import_time = time.perf_counter() - t_start
	logger.info(f"api_routes module imported successfully, time: {import_time:.3f}s")
	except Exception as e:
	import_time = time.perf_counter() - t_start
	logger.error(f"api_routes module import failed, time: {import_time:.3f}s, error: {e}")
	raise

	try:
	t_extract_start = time.perf_counter()
	extract_result = extract_chinese_celeb_dataset_sync()
	extract_time = time.perf_counter() - t_extract_start
	logger.info(
	"Chinese celeb dataset extracted successfully, time: %.3fs, target: %s",
	extract_time,
	extract_result.get("target_dir"),
	)
	except Exception as exc:
	logger.error(f"Failed to extract Chinese celeb dataset automatically: {exc}")
	raise


	@asynccontextmanager
	async def lifespan(app: FastAPI):
	start_time = time.perf_counter()
	logger.info("FaceScore service starting...")
	logger.info(f"Output directory: {OUTPUT_DIR}")
	logger.info(f"DeepFace available: {DEEPFACE_AVAILABLE}")
	logger.info(f"YOLO available: {YOLO_AVAILABLE}")
	logger.info(f"MediaPipe available: {DLIB_AVAILABLE}")
	logger.info(f"Archive directory: {IMAGES_DIR}")
	os.makedirs(OUTPUT_DIR, exist_ok=True)

	# 初始化数据库连接池
	try:
	await init_mysql_pool()
	logger.info("MySQL 连接池初始化完成")
	except Exception as exc:
	logger.error(f"初始化 MySQL 连接池失败: {exc}")
	raise

	# 启动图片清理定时任务
	logger.info("Starting image cleanup scheduled task...")
	try:
	start_cleanup_scheduler()
	logger.info("Image cleanup scheduled task started successfully")
	except Exception as e:
	logger.error(f"Failed to start image cleanup scheduled task: {e}")

	# 记录启动完成时间
	total_startup_time = time.perf_counter() - start_time
	logger.info(f"FaceScore service startup completed, total time: {total_startup_time:.3f}s")

	yield

	# 应用关闭时停止定时任务
	logger.info("Stopping image cleanup scheduled task...")
	try:
	stop_cleanup_scheduler()
	logger.info("Image cleanup scheduled task stopped")
	except Exception as e:
	logger.error(f"Failed to stop image cleanup scheduled task: {e}")

	# 关闭数据库连接池
	try:
	await close_mysql_pool()
	except Exception as exc:
	logger.warning(f"关闭 MySQL 连接池失败: {exc}")


	# 创建 FastAPI 应用
	app = FastAPI(
	title="Enhanced FaceScore 服务",
	description="支持多模型的人脸分析REST API服务，包含五官评分功能。支持混合模式：HowCuteAmI（颜值+性别）+ DeepFace（年龄+情绪）",
	version="3.0.0",
	docs_url="/cp_docs",
	redoc_url="/cp_redoc",
	lifespan=lifespan,
	)

	app.add_middleware(
	CORSMiddleware,
	allow_origins=["*"],
	allow_methods=["*"],
	allow_headers=["*"],
	)

	# 注册路由
	app.include_router(api_router)


	_default_frontend_dir = os.path.abspath(
	os.path.join(os.path.dirname(__file__), "facelist-web")
	)
	FRONTEND_DIR = os.path.abspath(
	os.path.expanduser(os.environ.get("FACELIST_FRONTEND_DIR", _default_frontend_dir))
	)

	_basic_user = os.environ.get("FACELIST_BASIC_USER", "admin")
	_basic_pass = os.environ.get("FACELIST_BASIC_PASS", "admin")
	_basic_secret = f"{_basic_user}:{_basic_pass}"
	_basic_token = "Basic " + base64.b64encode(_basic_secret.encode()).decode()


	class FacelistBasicAuthMiddleware(BaseHTTPMiddleware):
	"""仅保护 /facelist 前缀的 Basic Auth"""

	async def dispatch(self, request, call_next):
	path = request.url.path
	if path.startswith("/facelist"):
	auth = request.headers.get("Authorization", "")
	if not secrets.compare_digest(auth, _basic_token):
	return PlainTextResponse(
	"Unauthorized",
	status_code=401,
	headers={"WWW-Authenticate": 'Basic realm="facelist"'},
	)
	return await call_next(request)


	class SPAStaticFiles(StaticFiles):
	"""支持SPA路由回退到index.html的静态文件服务"""

	async def get_response(self, path: str, scope):
	response = await super().get_response(path, scope)
	if response.status_code == 404:
	# 未匹配到具体文件时回退到前端入口
	response = await super().get_response("index.html", scope)
	return response

	app.add_middleware(FacelistBasicAuthMiddleware)

	if os.path.isdir(FRONTEND_DIR):
	app.mount(
	"/facelist",
	SPAStaticFiles(directory=FRONTEND_DIR, html=True),
	name="facelist",
	)
	logger.info("前端静态资源已挂载在 /facelist ，目录: %s", FRONTEND_DIR)
	else:
	logger.warning("未找到前端目录 %s ，跳过前端静态资源挂载", FRONTEND_DIR)

	# 添加根路径处理
	@app.get("/")
	async def root():
	return "UP"


	if __name__ == "__main__":
	import uvicorn

	if not os.path.exists(MODELS_PATH):
	logger.critical(
	"Warning: 'models' directory not found. Please ensure it exists and contains model files."
	)
	logger.critical(
	"Exiting application as FaceAnalyzer cannot be initialized without models."
	)
	exit(1)

	# 根据日志开关配置 Uvicorn 日志
	if ENABLE_LOGGING:
	uvicorn.run(app, host="0.0.0.0", port=8080, reload=False)
	else:
	# 禁用 Uvicorn 的访问日志和错误日志
	uvicorn.run(
	app,
	host="0.0.0.0",
	port=8080,
	reload=False,
	access_log=False, # 禁用访问日志
	log_level="critical" # 只显示严重错误
	)