初始化提交

tianmu/src/main.py

@@ -0,0 +1,123 @@
+"""
+天目 · 魂镜本地 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()