File size: 5,778 Bytes
8da4eae 3b917ab 9f1fca5 902acbb 3b917ab 9f1fca5 3b917ab 9f1fca5 8da4eae 3b917ab 8da4eae 68c722b 3b917ab 68c722b 3b917ab 68c722b 8da4eae 3b917ab 8da4eae 3b917ab 8da4eae 3b917ab 8da4eae 3b917ab 8da4eae 3b917ab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | 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 """<!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)
# --- 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
|