parent
4de8e64451
commit
64cfbd851e
@ -0,0 +1,217 @@ |
||||
import os |
||||
import tempfile |
||||
import logging |
||||
import time |
||||
from fastapi import FastAPI, HTTPException, status, Body |
||||
from pydantic import BaseModel |
||||
import fitz # PyMuPDF |
||||
from paddleocr import PaddleOCR |
||||
from typing import Optional, Dict, Any |
||||
|
||||
# 日志配置保持不变 |
||||
logging.basicConfig( |
||||
level=logging.INFO, |
||||
format='%(asctime)s [%(levelname)s] %(message)s', |
||||
handlers=[ |
||||
logging.FileHandler('/var/log/ocr_service.log', encoding='utf-8'), |
||||
logging.StreamHandler() |
||||
] |
||||
) |
||||
logger = logging.getLogger('OCRService') |
||||
logger.setLevel(logging.DEBUG) |
||||
|
||||
# 安全配置保持不变 |
||||
ALLOWED_EXTENSIONS = {'pdf', 'jpg', 'jpeg', 'png'} |
||||
MAX_FILE_SIZE = 100 * 1024 * 1024 # 100MB |
||||
BASE_STORAGE = '/data/files' |
||||
|
||||
# OCR引擎配置保持不变 |
||||
ocr_engine = PaddleOCR( |
||||
use_angle_cls=True, |
||||
lang='ch', |
||||
enable_mkldnn=True, |
||||
det_limit_side_len=2200, |
||||
det_db_box_thresh=0.4 |
||||
) |
||||
|
||||
app = FastAPI() |
||||
|
||||
class OCRRequest(BaseModel): |
||||
file_path: str |
||||
|
||||
class OCRResponse(BaseModel): |
||||
status: str |
||||
data: Optional[Dict[str, Any]] = None |
||||
message: Optional[str] = None |
||||
code: Optional[int] = None |
||||
elapsed: str |
||||
|
||||
class OCRServiceError(Exception): |
||||
def __init__(self, message: str, status_code: int = 400): |
||||
self.message = message |
||||
self.status_code = status_code |
||||
|
||||
# 路径解析和验证函数保持不变 |
||||
def safe_resolve_path(user_path: str) -> str: |
||||
try: |
||||
clean_path = os.path.normpath(user_path).lstrip('/') |
||||
if not clean_path: |
||||
raise ValueError("空路径参数") |
||||
|
||||
abs_path = os.path.abspath(os.path.join(BASE_STORAGE, clean_path)) |
||||
|
||||
if not abs_path.startswith(BASE_STORAGE): |
||||
raise ValueError("非法路径访问") |
||||
|
||||
return abs_path |
||||
except Exception as e: |
||||
logger.error(f"路径解析失败: {str(e)}") |
||||
raise OCRServiceError("无效文件路径") from e |
||||
|
||||
def validate_file(file_path: str) -> str: |
||||
try: |
||||
if not os.path.exists(file_path): |
||||
raise OCRServiceError("文件不存在") |
||||
if not os.path.isfile(file_path): |
||||
raise OCRServiceError("路径不是文件") |
||||
if os.path.getsize(file_path) > MAX_FILE_SIZE: |
||||
raise OCRServiceError(f"文件超过大小限制 ({MAX_FILE_SIZE//1024//1024}MB)") |
||||
|
||||
_, ext_with_dot = os.path.splitext(file_path) |
||||
ext = ext_with_dot.lstrip('.').lower() |
||||
if ext not in ALLOWED_EXTENSIONS: |
||||
raise OCRServiceError(f"不支持的文件类型: {ext}") |
||||
|
||||
return ext |
||||
except OCRServiceError: |
||||
raise |
||||
except Exception as e: |
||||
logger.error(f"文件验证异常: {str(e)}") |
||||
raise OCRServiceError("文件验证失败") from e |
||||
|
||||
# 结果格式化和处理函数保持不变 |
||||
def format_ocr_result(raw_data): |
||||
results = [] |
||||
try: |
||||
for group in raw_data[0]: |
||||
boxes, (text, score) = group |
||||
logger.error(f"识别内容: {str(text)}") |
||||
results.append({ |
||||
"text": text, |
||||
"score": float(score), |
||||
"boxes": [list(map(float, point)) for point in boxes] |
||||
}) |
||||
return results |
||||
except Exception as e: |
||||
logger.error("解析失败：%s", str(e), exc_info=True) |
||||
return [] |
||||
|
||||
def process_pdf(file_path: str): |
||||
try: |
||||
logger.info(f"开始处理PDF: {file_path}") |
||||
doc = fitz.open(file_path) |
||||
if doc.is_encrypted: |
||||
if not doc.authenticate(""): |
||||
raise OCRServiceError("加密PDF需要密码") |
||||
|
||||
pages = [] |
||||
for page_num in range(len(doc)): |
||||
page_start = time.time() |
||||
page = doc.load_page(page_num) |
||||
zoom = 900 / 72 |
||||
mat = fitz.Matrix(zoom, zoom) |
||||
pix = page.get_pixmap( |
||||
matrix=mat, |
||||
colorspace=fitz.csGRAY, |
||||
alpha=False, |
||||
dpi=900 |
||||
) |
||||
with tempfile.NamedTemporaryFile(suffix=".png") as tmp: |
||||
pix.save(tmp.name) |
||||
page_result = ocr_engine.ocr(tmp.name, cls=True) |
||||
pages.append({ |
||||
"page": page_num + 1, |
||||
"content": format_ocr_result(page_result), |
||||
"process_time": f"{time.time() - page_start:.2f}s" |
||||
}) |
||||
|
||||
return { |
||||
"type": "pdf", |
||||
"page_count": len(doc), |
||||
"pages": pages |
||||
} |
||||
except OCRServiceError: |
||||
raise |
||||
except Exception as e: |
||||
logger.error(f"PDF处理异常: {str(e)}") |
||||
raise OCRServiceError("PDF处理失败") from e |
||||
|
||||
def process_image(file_path: str): |
||||
try: |
||||
logger.info(f"开始处理图像: {file_path}") |
||||
start_time = time.time() |
||||
result = ocr_engine.ocr(file_path, cls=True) |
||||
logger.debug("原始OCR数据结构类型: %s", type(result)) |
||||
formatted = format_ocr_result(result) |
||||
return { |
||||
"type": "image", |
||||
"results": formatted |
||||
} |
||||
except Exception as e: |
||||
logger.error("图像处理异常") |
||||
raise |
||||
|
||||
@app.exception_handler(OCRServiceError) |
||||
async def ocr_exception_handler(request, exc: OCRServiceError): |
||||
return JSONResponse( |
||||
status_code=exc.status_code, |
||||
content={ |
||||
"status": "error", |
||||
"message": exc.message, |
||||
"code": exc.status_code, |
||||
"elapsed": "0.00s" |
||||
} |
||||
) |
||||
|
||||
@app.post("/ocr", response_model=OCRResponse) |
||||
async def ocr_service(request_data: OCRRequest = Body(...)): |
||||
start_time = time.time() |
||||
response = {"status": "success", "data": None, "message": None, "code": None} |
||||
|
||||
try: |
||||
abs_path = safe_resolve_path(request_data.file_path) |
||||
logger.debug(f"处理请求文件: {abs_path}") |
||||
ext = validate_file(abs_path) |
||||
|
||||
if ext == 'pdf': |
||||
result = process_pdf(abs_path) |
||||
else: |
||||
result = process_image(abs_path) |
||||
|
||||
response["data"] = result |
||||
|
||||
except OCRServiceError as e: |
||||
raise e |
||||
except Exception as e: |
||||
logger.error(f"系统异常: {str(e)}", exc_info=True) |
||||
raise HTTPException( |
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, |
||||
detail="系统内部错误" |
||||
) |
||||
finally: |
||||
response["elapsed"] = f"{time.time() - start_time:.2f}s" |
||||
|
||||
return response |
||||
|
||||
@app.get("/healthcheck", response_model=Dict[str, Any]) |
||||
async def health_check(): |
||||
return { |
||||
"status": "ok", |
||||
"timestamp": time.time(), |
||||
"service": "OCR" |
||||
} |
||||
|
||||
if __name__ == "__main__": |
||||
import uvicorn |
||||
os.makedirs(BASE_STORAGE, exist_ok=True) |
||||
uvicorn.run(app, host="0.0.0.0", port=5000) |
||||
Loading…
Reference in new issue