"""
OCR Internal Service
提供标准整图 OCR 识别服务。
"""
import io
import asyncio
from typing import List, Optional
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
from PIL import Image
import numpy as np
from datetime import datetime
import logging
from config.config import settings
from engine.ocr_engine import OCRProcessor
# 配置日志
logging.basicConfig(
level=logging.DEBUG if settings.debug else logging.INFO,
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
)
logger = logging.getLogger(__name__)
app = FastAPI(title="OCR Internal Service", version="2.0.0")
# 初始化OCR处理器
ocr = OCRProcessor(
use_det=settings.ocr_use_det,
use_cls=settings.ocr_use_cls,
use_rec=settings.ocr_use_rec,
model_dir=settings.ocr_model_dir,
text_score=settings.ocr_text_score,
max_side_len=settings.ocr_max_side_len,
det_limit_side_len=settings.ocr_det_limit_side_len,
det_box_thresh=settings.ocr_det_box_thresh,
det_thresh=settings.ocr_det_thresh,
rec_batch_num=settings.ocr_rec_batch_num,
device=settings.ocr_device,
device_id=settings.ocr_device_id,
intra_threads=settings.ocr_intra_threads,
inter_threads=settings.ocr_inter_threads,
enable_vertical=settings.ocr_enable_vertical,
)
# ─── 响应模型 ─────────────────────────────────────────────
class TextRegion(BaseModel):
confidence: float
text: str
text_region: List[List[int]] # [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
class OCRResponse(BaseModel):
code: int
message: str
data: Optional[List[TextRegion]] = None
timestamp: str
# ─── 辅助函数 ─────────────────────────────────────────────
def _build_response(word_results) -> OCRResponse:
"""将 word_results 构建为响应，兼容嵌套和扁平两种格式"""
ocr_result = []
if not word_results:
return OCRResponse(
code=0, message="success", data=ocr_result,
timestamp=datetime.now().isoformat(),
)
for line in word_results:
if isinstance(line, (tuple, list)) and len(line) >= 3 and isinstance(line[0], str):
# 扁平 item: 单条 (text, conf, coords)，line 本身就是 item
text, confidence, coords = line[0], line[1], line[2]
if coords is not None:
ocr_result.append(TextRegion(
confidence=confidence, text=text, text_region=coords,
))
elif isinstance(line, (tuple, list)):
# 嵌套结构: line = ((item1), (item2), ...)
for item in line:
if isinstance(item, (tuple, list)) and len(item) >= 3:
text, confidence, coords = item[0], item[1], item[2]
if coords is not None:
ocr_result.append(TextRegion(
confidence=confidence, text=text, text_region=coords,
))
return OCRResponse(
code=0,
message="success",
data=ocr_result,
timestamp=datetime.now().isoformat(),
)
def _error_response(msg: str) -> OCRResponse:
return OCRResponse(
code=500,
message=msg,
timestamp=datetime.now().isoformat(),
)
# ─── API 端点 ─────────────────────────────────────────────
@app.get("/ocr_system/health")
async def health_check():
"""健康检查"""
return {
"status": "ok",
"service": "ocr",
"version": "2.0.0",
"device": ocr.device,
"timestamp": datetime.now().isoformat(),
}
@app.post("/ocr_system/ocr", response_model=OCRResponse)
async def ocr_recognition(file: UploadFile = File(...)):
"""上传图片文件并执行整图 OCR 识别。"""
try:
# 读取文件
image_bytes = await file.read()
image_file = io.BytesIO(image_bytes)
image = Image.open(image_file)
image_np = np.array(image)
logger.info(f"收到OCR请求: size={image_np.shape}")
# 预处理（原图直通：灰度化+CLAHE+锐化会损失边缘小字细节，已实测降低召回）
loop = asyncio.get_event_loop()
im_np = image_np
logger.info("使用标准OCR模式")
result = await loop.run_in_executor(None, ocr.process_with_vertical, im_np)
if result is None:
return _error_response("OCR处理失败")
return _build_response(result.word_results)
except Exception as e:
logger.error(f"OCR处理异常: {str(e)}", exc_info=True)
return _error_response(f"处理异常: {str(e)}")
# ─── 可视化页面 ───────────────────────────────────────────
@app.get("/ocr_system/viewer", response_class=HTMLResponse)
async def viewer():
"""OCR结果可视化页面"""
return HTMLResponse(content=VIEWER_HTML)
VIEWER_HTML = r"""
OCR 可视化工具
上传图片
+
点击或拖拽上传
JPG / PNG / BMP / TIFF
"""
# ─── 启动入口 ─────────────────────────────────────────────
if __name__ == "__main__":
import uvicorn
logger.info(f"启动OCR服务: {settings.host}:{settings.port}")
uvicorn.run(
"main:app",
host=settings.host,
port=settings.port,
reload=settings.debug
)