124 lines
3.4 KiB
Python
124 lines
3.4 KiB
Python
"""
|
||
天目 · 魂镜本地 OCR 引擎 — FastAPI 应用入口。
|
||
|
||
提供通用文字识别（PP-OCRv6，via RapidOCR）和复杂文档解析（PaddleOCR-VL）API。
|
||
所有模型由本地引擎直接加载推理，不依赖外部服务。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
|
||
# 抑制进程退出时 multiprocessing 信号量泄漏警告（强制覆盖，子进程继承）
|
||
os.environ["PYTHONWARNINGS"] = "ignore::UserWarning"
|
||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||
|
||
import gc
|
||
import warnings
|
||
warnings.filterwarnings("ignore", message="resource_tracker:")
|
||
|
||
from fastapi import FastAPI
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
|
||
from src import __version__
|
||
from src.config import settings
|
||
from src.engine.rapidocr import RapidOCREngine
|
||
from src.engine.paddleocr_vl import PaddleOCRVLEngine
|
||
from src.endpoints import health, ocr
|
||
from src.utils.logging import logger, setup_logging
|
||
|
||
setup_logging(settings.log_level)
|
||
|
||
app = FastAPI(
|
||
title="天目 · 魂镜本地 OCR 引擎",
|
||
version=__version__,
|
||
description="本地 OCR 推理引擎 — PP-OCRv6（快速文字提取）+ PaddleOCR-VL（复杂文档解析）",
|
||
)
|
||
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
# 全局引擎实例
|
||
rapidocr_engine = RapidOCREngine()
|
||
paddleocr_vl_engine = PaddleOCRVLEngine()
|
||
|
||
|
||
@app.on_event("startup")
|
||
async def startup() -> None:
|
||
logger.info(
|
||
"天目启动 | host={} port={} device={}(detected={}) models_dir={}",
|
||
settings.host,
|
||
settings.port,
|
||
settings.device,
|
||
settings.detected_device,
|
||
settings.models_dir_resolved,
|
||
)
|
||
# 注入引擎实例到端点模块
|
||
health.init(rapidocr_engine, paddleocr_vl_engine)
|
||
ocr.init(rapidocr_engine, paddleocr_vl_engine)
|
||
|
||
# 懒加载 RapidOCR 引擎（模型存在则加载，失败不阻塞启动）
|
||
if settings.ppocrv6_dir.exists():
|
||
try:
|
||
await rapidocr_engine.load()
|
||
except Exception as e:
|
||
logger.warning("PP-OCRv6 引擎加载失败（可在运行时重试）: {}", e)
|
||
else:
|
||
logger.warning("PP-OCRv6 模型目录不存在: {}", settings.ppocrv6_dir)
|
||
|
||
|
||
@app.on_event("shutdown")
|
||
async def shutdown() -> None:
|
||
"""优雅关闭：释放引擎资源，超时兜底。"""
|
||
import asyncio
|
||
|
||
try:
|
||
await asyncio.wait_for(
|
||
asyncio.gather(
|
||
rapidocr_engine.close(),
|
||
paddleocr_vl_engine.close(),
|
||
),
|
||
timeout=5.0,
|
||
)
|
||
except asyncio.TimeoutError:
|
||
logger.warning("引擎关闭超时，强制退出")
|
||
except Exception as e:
|
||
logger.warning("引擎关闭异常: {}", e)
|
||
|
||
gc.collect()
|
||
logger.info("天目停止")
|
||
|
||
|
||
# 注册路由
|
||
app.include_router(health.router, tags=["health"])
|
||
app.include_router(ocr.router, tags=["ocr"])
|
||
|
||
|
||
def _run() -> None:
|
||
"""启动 uvicorn 服务。"""
|
||
import uvicorn
|
||
uvicorn.run(
|
||
"src.main:app",
|
||
host=settings.host,
|
||
port=settings.port,
|
||
log_level=settings.log_level.lower(),
|
||
)
|
||
|
||
|
||
def main_cli() -> None:
|
||
"""uv 命令行入口（uv run tianmu），确保 PYTHONWARNINGS 在 reloader 之前设置。"""
|
||
import sys
|
||
os.environ["PYTHONWARNINGS"] = "ignore::UserWarning"
|
||
sys.argv = [a for a in sys.argv if a != "tianmu"]
|
||
_run()
|
||
|
||
|
||
# 直接运行入口（python -m src.main）
|
||
if __name__ == "__main__":
|
||
_run()
|