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 # 2 phiên bản model chạy song song trong CÙNG 1 Space để so sánh. # Mỗi version tải từ 1 HF model repo riêng vào thư mục riêng. 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 """ VOCR — Vietnamese OCR

📖 VOCR — OCR tiếng Việt sách cũ

PP-OCRv5 fine-tune (det hmean 0.916) + hậu xử lý ViLM. 2 phiên bản rec để so sánh:

So sánh trên cùng ảnh:

curl -X POST {URL}/v4/ocr  -F "file=@trang.png"
curl -X POST {URL}/v11/ocr -F "file=@trang.png"
""" @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) # --- v4 (bản cũ) --- @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 # --- v11 (bản mới) --- @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 # --- /ocr mặc định = v11 (giữ tương thích client cũ) --- @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