| import os |
| import tempfile |
| from pathlib import Path |
| from contextlib import asynccontextmanager |
| from typing import List, Optional |
|
|
| from fastapi import FastAPI, File, UploadFile, HTTPException |
| from fastapi.responses import PlainTextResponse |
| from pydantic import BaseModel |
|
|
| from vocr_pipeline import VOCRPipeline |
|
|
|
|
| |
| |
| MODEL_REPOS = { |
| "v4": os.getenv("HF_MODEL_REPO_V4", "sirimiri/vocr-models"), |
| "v11": os.getenv("HF_MODEL_REPO_V11", "sirimiri/vocr-models-v11"), |
| } |
| pipelines: dict = {} |
|
|
|
|
| def _download_repo(repo: str, target: str): |
| """Tải 1 model repo về `target/` (giữ structure models/ dict/ lm_data/). Bỏ qua nếu đã có.""" |
| if os.path.exists(os.path.join(target, "models", "inference_rec")): |
| print(f"[{repo}] đã có ở {target}, bỏ qua tải.") |
| return |
| try: |
| from huggingface_hub import snapshot_download |
| hf_token = os.getenv("HF_TOKEN", None) |
| print(f"Downloading {repo} -> {target} ...") |
| snapshot_download(repo_id=repo, local_dir=target, token=hf_token) |
| print(f"[{repo}] downloaded.") |
| except Exception as e: |
| print(f"[{repo}] download failed: {e}") |
|
|
|
|
| @asynccontextmanager |
| async def lifespan(app: FastAPI): |
| beam = int(os.getenv("VOCR_BEAM_SIZE", "0")) |
| no_enh = os.getenv("VOCR_NO_ENHANCE", "1") == "1" |
| for ver, repo in MODEL_REPOS.items(): |
| target = f"./hf_{ver}" |
| _download_repo(repo, target) |
| try: |
| pipelines[ver] = VOCRPipeline( |
| model_dir=f"{target}/models", |
| dict_dir=f"{target}/dict", |
| device="cpu", |
| output_dir=f"./output/api_{ver}", |
| beam_size=beam, |
| no_enhance=no_enh, |
| ) |
| print(f"[{ver}] pipeline sẵn sàng (repo {repo}).") |
| except Exception as e: |
| print(f"[{ver}] init pipeline failed: {e}") |
| yield |
|
|
| app = FastAPI(title="VOCR API", version="1.0.0", lifespan=lifespan) |
|
|
|
|
| class RegionResult(BaseModel): |
| box: List[List[float]] |
| text: str |
| confidence: float |
|
|
| class OCRResponse(BaseModel): |
| full_text: str |
| avg_confidence: float |
| detections: int |
| results: List[RegionResult] |
|
|
|
|
| from fastapi.responses import HTMLResponse |
|
|
| @app.get("/", response_class=HTMLResponse) |
| def home(): |
| return """<!doctype html><html lang=vi><head><meta charset=utf-8> |
| <title>VOCR — Vietnamese OCR</title> |
| <style>body{font-family:system-ui;max-width:640px;margin:60px auto;padding:0 20px;line-height:1.6} |
| code,pre{background:#f4f4f0;padding:2px 6px;border-radius:4px}pre{padding:12px;overflow-x:auto}</style></head> |
| <body><h1>📖 VOCR — OCR tiếng Việt sách cũ</h1> |
| <p>PP-OCRv5 fine-tune (det hmean 0.916) + hậu xử lý ViLM. <b>2 phiên bản rec để so sánh:</b></p> |
| <ul> |
| <li><b>v4</b> (cũ) — <code>/v4/ocr</code></li> |
| <li><b>v11</b> (mới: acc 0.225, giảm ~27% lỗi lặp ký tự) — <code>/v11/ocr</code></li> |
| <li><code>/ocr</code> = mặc định trỏ v11</li> |
| </ul> |
| <ul><li><a href="/docs">Swagger UI (/docs)</a></li><li><a href="/health">/health</a></li></ul> |
| <p>So sánh trên cùng ảnh:</p> |
| <pre>curl -X POST {URL}/v4/ocr -F "file=@trang.png" |
| curl -X POST {URL}/v11/ocr -F "file=@trang.png"</pre> |
| </body></html>""" |
|
|
| @app.get("/health") |
| def health(): |
| return { |
| "status": "ok", |
| "versions": {ver: (ver in pipelines) for ver in MODEL_REPOS}, |
| "repos": MODEL_REPOS, |
| } |
|
|
|
|
| async def _run_ocr(ver: str, file: UploadFile) -> OCRResponse: |
| pipe = pipelines.get(ver) |
| if pipe is None: |
| raise HTTPException(status_code=503, detail=f"Model {ver} chưa sẵn sàng") |
|
|
| ext = Path(file.filename).suffix.lower() |
| if ext not in {".jpg", ".jpeg", ".png", ".bmp", ".tiff"}: |
| raise HTTPException(status_code=400, detail=f"Định dạng không hỗ trợ: {ext}") |
|
|
| with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp: |
| tmp.write(await file.read()) |
| tmp_path = tmp.name |
|
|
| try: |
| result = pipe.run(tmp_path, save=False) |
| if "error" in result: |
| raise HTTPException(status_code=422, detail=result["error"]) |
| return OCRResponse( |
| full_text=result["full_text"], |
| avg_confidence=round(result["avg_confidence"], 4), |
| detections=result["detections"], |
| results=[ |
| RegionResult( |
| box=r["box"], |
| text=r["best_text"], |
| confidence=round(r["best_score"], 4), |
| ) |
| for r in result["results"] |
| ], |
| ) |
| finally: |
| os.remove(tmp_path) |
|
|
|
|
| |
| @app.post("/v4/ocr", response_model=OCRResponse) |
| async def ocr_v4(file: UploadFile = File(...)): |
| return await _run_ocr("v4", file) |
|
|
| @app.post("/v4/ocr/text", response_class=PlainTextResponse) |
| async def ocr_v4_text(file: UploadFile = File(...)): |
| return (await _run_ocr("v4", file)).full_text |
|
|
|
|
| |
| @app.post("/v11/ocr", response_model=OCRResponse) |
| async def ocr_v11(file: UploadFile = File(...)): |
| return await _run_ocr("v11", file) |
|
|
| @app.post("/v11/ocr/text", response_class=PlainTextResponse) |
| async def ocr_v11_text(file: UploadFile = File(...)): |
| return (await _run_ocr("v11", file)).full_text |
|
|
|
|
| |
| @app.post("/ocr", response_model=OCRResponse) |
| async def ocr(file: UploadFile = File(...)): |
| return await _run_ocr("v11", file) |
|
|
| @app.post("/ocr/text", response_class=PlainTextResponse) |
| async def ocr_text(file: UploadFile = File(...)): |
| return (await _run_ocr("v11", file)).full_text |
|
|
