File size: 5,757 Bytes
9e7ba62 | 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 | """
app.py - FastAPI server: serves the UI at / and the OCR API at /api/*.
uvicorn app:app --host 0.0.0.0 --port 8000 # real model (GPU)
OCR_MOCK=1 uvicorn app:app --port 8000 # UI development, no GPU
"""
import difflib, os, re, shutil, tempfile, threading, time, uuid
from pathlib import Path
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from ocr_engine import CONFIG, IMAGE_EXT, LANG_NAMES, OCREngine, mask_ids
try: # benchmark.py needs validate.py; without it, no ground-truth scoring
from benchmark import score
except ImportError:
score = None
MAX_UPLOAD_MB = int(os.getenv("OCR_MAX_UPLOAD_MB", "25"))
JOB_TTL_S = int(os.getenv("OCR_JOB_TTL_S", "3600")) # uploads + outputs deleted after this
JOBS_DIR = Path(os.getenv("OCR_JOBS_DIR", tempfile.gettempdir())) / "ocr_jobs"
JOBS_DIR.mkdir(parents=True, exist_ok=True)
STATIC = Path(__file__).parent / "static"
app = FastAPI(title="Indic Handwritten OCR")
ENGINE: OCREngine | None = None
_load_error: str | None = None
def get_engine() -> OCREngine:
"""Loads the model once. Modal / Colab can also inject a pre-loaded engine via set_engine()."""
global ENGINE, _load_error
if ENGINE is None:
try:
ENGINE = OCREngine.load()
except Exception as e: # surface load failures to /api/health
_load_error = repr(e)
raise
return ENGINE
def set_engine(engine: OCREngine):
global ENGINE
ENGINE = engine
@app.on_event("startup")
def _warm():
if os.getenv("OCR_EAGER_LOAD", "1") == "1":
threading.Thread(target=get_engine, daemon=True).start()
threading.Thread(target=_janitor, daemon=True).start()
def _janitor():
while True:
cutoff = time.time() - JOB_TTL_S
for d in JOBS_DIR.iterdir():
if d.is_dir() and d.stat().st_mtime < cutoff:
shutil.rmtree(d, ignore_errors=True)
time.sleep(300)
@app.get("/api/health")
def health():
return {"ready": bool(ENGINE and ENGINE.ready), "mock": CONFIG["MOCK"], "error": _load_error,
"languages": LANG_NAMES, "max_upload_mb": MAX_UPLOAD_MB, "max_pages": CONFIG["MAX_PAGES"],
"scoring": score is not None}
def evaluate(result: dict, gt: str, lang: str | None) -> dict:
"""Same metrics as benchmark.py's run_docs(), plus a word diff for the UI."""
pred = "\n".join(b["text"] for p in result["pages"] for b in p["blocks"])
gt = mask_ids(gt.strip()) # the output is masked, so mask the ground truth the same way
tok = lambda s: re.findall(r"\S+", s)
g, p = tok(gt), tok(pred)
diff = []
for op, i0, i1, j0, j1 in difflib.SequenceMatcher(None, g, p, autojunk=False).get_opcodes():
if op == "equal":
diff.append(["eq", " ".join(g[i0:i1])])
if op in ("delete", "replace"):
diff.append(["del", " ".join(g[i0:i1])]) # in ground truth, missing from OCR
if op in ("insert", "replace"):
diff.append(["ins", " ".join(p[j0:j1])]) # in OCR, not in ground truth
return {**score(gt, pred, lang or "-"), "gt_chars": len(gt), "gt_words": len(g), "diff": diff}
@app.post("/api/ocr")
def ocr(file: UploadFile = File(...), lang: str = Form(""), gt: UploadFile | None = File(None)):
# sync def -> FastAPI runs it in a worker thread; the engine lock serialises GPU access
suffix = Path(file.filename or "").suffix.lower()
if suffix not in IMAGE_EXT | {".pdf"}:
raise HTTPException(400, f"Unsupported file type '{suffix}'. Use PDF or an image.")
job_id = uuid.uuid4().hex[:12]
job_dir = JOBS_DIR / job_id
job_dir.mkdir()
safe_name = Path(file.filename).name.replace(" ", "_")
src = job_dir / safe_name
size = 0
with src.open("wb") as fh:
while chunk := file.file.read(1 << 20):
size += len(chunk)
if size > MAX_UPLOAD_MB << 20:
shutil.rmtree(job_dir, ignore_errors=True)
raise HTTPException(413, f"File is larger than {MAX_UPLOAD_MB} MB")
fh.write(chunk)
gt_text = None
if gt is not None and gt.filename:
if score is None:
shutil.rmtree(job_dir, ignore_errors=True)
raise HTTPException(400, "Ground-truth scoring needs validate.py and benchmark.py on the server")
try:
gt_text = gt.file.read(5 << 20).decode("utf-8-sig")
except UnicodeDecodeError:
shutil.rmtree(job_dir, ignore_errors=True)
raise HTTPException(400, "Ground truth must be a UTF-8 .txt file")
try:
result = get_engine().process(src, lang or None, job_dir / "out")
except ValueError as e:
shutil.rmtree(job_dir, ignore_errors=True)
raise HTTPException(400, str(e))
except Exception as e:
shutil.rmtree(job_dir, ignore_errors=True)
raise HTTPException(500, f"OCR failed: {e}")
src.unlink(missing_ok=True) # don't keep the uploaded document around
result["job_id"] = job_id
if gt_text is not None:
result["eval"] = evaluate(result, gt_text, lang or None)
return JSONResponse(result)
@app.get("/api/jobs/{job_id}/{name}")
def download(job_id: str, name: str):
if not job_id.isalnum() or "/" in name or ".." in name:
raise HTTPException(400, "bad path")
p = JOBS_DIR / job_id / "out" / name
if not p.is_file():
raise HTTPException(404, "expired or not found")
return FileResponse(p, filename=name)
app.mount("/", StaticFiles(directory=STATIC, html=True), name="ui")
|