| """ |
| 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: |
| 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")) |
| 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: |
| _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()) |
| 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])]) |
| if op in ("insert", "replace"): |
| diff.append(["ins", " ".join(p[j0:j1])]) |
| 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)): |
| |
| 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) |
| 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") |
|
|
