Spaces:

Ratnesh-dev
/

pdf-ocr

Sleeping

File size: 11,601 Bytes

import os
import shutil
import subprocess
import tempfile
import time
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
from typing import Any

import cv2
import numpy as np
import pypdfium2 as pdfium
from fastapi import FastAPI, File, HTTPException, Query, UploadFile
from fastapi.responses import HTMLResponse
from PIL import Image, UnidentifiedImageError

os.environ.setdefault("OMP_THREAD_LIMIT", "1")


def _read_cgroup_file(path: str) -> str | None:
    try:
        return Path(path).read_text().strip()
    except OSError:
        return None


def _parse_cpuset(cpu_set: str | None) -> int | None:
    if not cpu_set:
        return None
    count = 0
    for part in cpu_set.split(","):
        if "-" in part:
            start, end = part.split("-", 1)
            count += int(end) - int(start) + 1
        elif part.strip():
            count += 1
    return count or None


def _get_cpu_quota_count() -> int | None:
    cpu_max = _read_cgroup_file("/sys/fs/cgroup/cpu.max")
    if cpu_max:
        quota, period = cpu_max.split()[:2]
        if quota != "max":
            return max(1, int(int(quota) / int(period)))

    quota = _read_cgroup_file("/sys/fs/cgroup/cpu/cpu.cfs_quota_us")
    period = _read_cgroup_file("/sys/fs/cgroup/cpu/cpu.cfs_period_us")
    if quota and period and int(quota) > 0:
        return max(1, int(int(quota) / int(period)))
    return None


def _get_cpuset_count() -> int | None:
    return _parse_cpuset(
        _read_cgroup_file("/sys/fs/cgroup/cpuset.cpus.effective")
        or _read_cgroup_file("/sys/fs/cgroup/cpuset/cpuset.cpus")
    )


def _get_available_cpu_count() -> int:
    candidates = [os.cpu_count() or 1]
    if quota_count := _get_cpu_quota_count():
        candidates.append(quota_count)
    if cpuset_count := _get_cpuset_count():
        candidates.append(cpuset_count)
    return max(1, min(candidates))


def _get_cpu_details() -> dict[str, Any]:
    return {
        "host_cpu_count": os.cpu_count(),
        "quota_cpu_count": _get_cpu_quota_count(),
        "cpuset_cpu_count": _get_cpuset_count(),
        "available_cpu_count": _get_available_cpu_count(),
        "cpu_max": _read_cgroup_file("/sys/fs/cgroup/cpu.max"),
        "cpuset_cpus_effective": _read_cgroup_file("/sys/fs/cgroup/cpuset.cpus.effective"),
    }


DEFAULT_DPI = int(os.getenv("OCR_DPI", "220"))
DEFAULT_LANG = os.getenv("OCR_LANG", "eng")
DEFAULT_PSM = int(os.getenv("OCR_PSM", "6"))
MAX_UPLOAD_MB = int(os.getenv("MAX_UPLOAD_MB", "300"))
MAX_WORKERS = int(os.getenv("OCR_MAX_WORKERS", str(_get_available_cpu_count())))
PAGE_TIMEOUT_SECONDS = int(os.getenv("OCR_PAGE_TIMEOUT_SECONDS", "120"))

app = FastAPI(
    title="Screener OCR API",
    description="CPU-only OCR service returning pages_text arrays for scanned PDFs.",
    version="0.1.0",
)


@app.get("/", response_class=HTMLResponse)
def index() -> str:
    return """
    <html>
      <body>
        <h1>Screener OCR API</h1>
        <p>POST a PDF to <code>/ocr</code> as multipart field <code>file</code>.</p>
        <pre>curl -F "file=@sample.pdf" https://YOUR_SPACE.hf.space/ocr</pre>
        <p>Response shape: <code>{"pages_text": ["page 1", "page 2"]}</code></p>
      </body>
    </html>
    """


@app.get("/health")
def health() -> dict[str, Any]:
    return {
        "ok": True,
        "engine": "tesseract",
        **_get_cpu_details(),
        "max_workers": MAX_WORKERS,
        "tesseract_thread_limit": os.getenv("OMP_THREAD_LIMIT"),
        "tesseract_version": _tesseract_version(),
    }


@app.post("/ocr")
async def ocr(
    file: UploadFile = File(...),
    dpi: int = Query(DEFAULT_DPI, ge=150, le=350),
    lang: str = Query(DEFAULT_LANG, min_length=3, max_length=32),
    psm: int = Query(DEFAULT_PSM, ge=3, le=13),
    deskew: bool = Query(True),
    workers: int | None = Query(None, ge=1, le=MAX_WORKERS),
    first_page: int = Query(1, ge=1),
    last_page: int | None = Query(None, ge=1),
) -> dict[str, Any]:
    started_at = time.monotonic()
    suffix = Path(file.filename or "upload.pdf").suffix.lower() or ".pdf"

    with tempfile.TemporaryDirectory(prefix="ocr-api-") as temp_dir:
        input_path = Path(temp_dir) / f"input{suffix}"
        await _save_upload(file, input_path)

        try:
            if suffix in {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".webp"}:
                pages_text = [_ocr_image_file(str(input_path), dpi, lang, psm, deskew)]
                errors: dict[str, str] = {}
            else:
                pages_text, errors = _ocr_pdf(
                    str(input_path),
                    dpi=dpi,
                    lang=lang,
                    psm=psm,
                    deskew=deskew,
                    workers=workers or MAX_WORKERS,
                    first_page=first_page,
                    last_page=last_page,
                )
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc
        except Exception as exc:
            raise HTTPException(status_code=500, detail=f"OCR failed: {exc}") from exc

    response: dict[str, Any] = {
        "pages_text": pages_text,
        "page_count": len(pages_text),
        "processing_ms": round((time.monotonic() - started_at) * 1000),
    }
    if errors:
        response["errors"] = errors
    return response


async def _save_upload(file: UploadFile, destination: Path) -> None:
    max_bytes = MAX_UPLOAD_MB * 1024 * 1024
    size = 0
    with destination.open("wb") as output:
        while chunk := await file.read(1024 * 1024):
            size += len(chunk)
            if size > max_bytes:
                raise HTTPException(
                    status_code=413,
                    detail=f"File is larger than MAX_UPLOAD_MB={MAX_UPLOAD_MB}",
                )
            output.write(chunk)


def _ocr_pdf(
    pdf_path: str,
    *,
    dpi: int,
    lang: str,
    psm: int,
    deskew: bool,
    workers: int,
    first_page: int,
    last_page: int | None,
) -> tuple[list[str], dict[str, str]]:
    page_count = _get_pdf_page_count(pdf_path)
    if page_count == 0:
        return [], {}

    start_index = first_page - 1
    end_index = page_count if last_page is None else min(last_page, page_count)
    if start_index >= page_count:
        raise ValueError(f"first_page={first_page} exceeds PDF page count {page_count}")
    if end_index <= start_index:
        raise ValueError("last_page must be >= first_page")

    tasks = [
        (pdf_path, page_index, dpi, lang, psm, deskew)
        for page_index in range(start_index, end_index)
    ]
    worker_count = max(1, min(workers, len(tasks)))

    pages_text: list[str] = []
    errors: dict[str, str] = {}
    if worker_count == 1:
        results = [_ocr_pdf_page(task) for task in tasks]
    else:
        with ProcessPoolExecutor(max_workers=worker_count) as executor:
            results = list(executor.map(_ocr_pdf_page, tasks, chunksize=1))

    for page_number, text, error in results:
        pages_text.append(text)
        if error:
            errors[str(page_number)] = error
    return pages_text, errors


def _get_pdf_page_count(pdf_path: str) -> int:
    try:
        doc = pdfium.PdfDocument(pdf_path)
    except Exception as exc:
        raise ValueError("Uploaded file is not a readable PDF") from exc
    try:
        return len(doc)
    finally:
        doc.close()


def _ocr_pdf_page(task: tuple[str, int, int, str, int, bool]) -> tuple[int, str, str | None]:
    pdf_path, page_index, dpi, lang, psm, deskew = task
    page_number = page_index + 1
    try:
        image = _render_pdf_page(pdf_path, page_index, dpi)
        text = _ocr_image(image, dpi, lang, psm, deskew)
        return page_number, _clean_text(text), None
    except Exception as exc:
        return page_number, "", str(exc)[:500]


def _render_pdf_page(pdf_path: str, page_index: int, dpi: int) -> Image.Image:
    doc = pdfium.PdfDocument(pdf_path)
    try:
        page = doc[page_index]
        try:
            bitmap = page.render(scale=dpi / 72)
            return bitmap.to_pil().convert("RGB")
        finally:
            page.close()
    finally:
        doc.close()


def _ocr_image_file(image_path: str, dpi: int, lang: str, psm: int, deskew: bool) -> str:
    try:
        with Image.open(image_path) as image:
            return _clean_text(_ocr_image(image.convert("RGB"), dpi, lang, psm, deskew))
    except UnidentifiedImageError as exc:
        raise ValueError("Uploaded file is neither a readable PDF nor a supported image") from exc


def _ocr_image(image: Image.Image, dpi: int, lang: str, psm: int, deskew: bool) -> str:
    processed = _preprocess_for_ocr(image, deskew=deskew)
    with tempfile.NamedTemporaryFile(suffix=".png") as image_file:
        processed.save(image_file.name, format="PNG")
        cmd = [
            "tesseract",
            image_file.name,
            "stdout",
            "-l",
            lang,
            "--psm",
            str(psm),
            "--dpi",
            str(dpi),
            "-c",
            "preserve_interword_spaces=1",
        ]
        completed = subprocess.run(
            cmd,
            check=False,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=PAGE_TIMEOUT_SECONDS,
        )
    if completed.returncode != 0:
        raise RuntimeError(completed.stderr.strip() or "tesseract failed")
    return completed.stdout


def _preprocess_for_ocr(image: Image.Image, *, deskew: bool) -> Image.Image:
    rgb = np.array(image.convert("RGB"))
    if deskew:
        angle = _estimate_skew_angle(rgb)
        if 0.2 <= abs(angle) <= 15:
            rgb = _rotate_image(rgb, angle)
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
    return Image.fromarray(gray)


def _estimate_skew_angle(rgb: np.ndarray) -> float:
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
    height, width = gray.shape[:2]
    scale = min(1.0, 1800 / max(height, width))
    if scale < 1.0:
        gray = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)

    inverted = cv2.bitwise_not(gray)
    thresholded = cv2.threshold(
        inverted, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU
    )[1]
    coords = np.column_stack(np.where(thresholded > 0))
    if len(coords) < 100:
        return 0.0

    angle = cv2.minAreaRect(coords)[-1]
    if angle < -45:
        angle = -(90 + angle)
    else:
        angle = -angle
    return float(angle)


def _rotate_image(rgb: np.ndarray, angle: float) -> np.ndarray:
    height, width = rgb.shape[:2]
    center = (width / 2, height / 2)
    matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
    return cv2.warpAffine(
        rgb,
        matrix,
        (width, height),
        flags=cv2.INTER_CUBIC,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=(255, 255, 255),
    )


def _clean_text(text: str) -> str:
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    lines = [line.rstrip() for line in text.split("\n")]
    while lines and not lines[0].strip():
        lines.pop(0)
    while lines and not lines[-1].strip():
        lines.pop()
    return "\n".join(lines)


def _tesseract_version() -> str:
    executable = shutil.which("tesseract")
    if not executable:
        return "missing"
    completed = subprocess.run(
        [executable, "--version"],
        check=False,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        timeout=5,
    )
    return completed.stdout.splitlines()[0] if completed.stdout else "unknown"