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 """
Screener OCR API
POST a PDF to /ocr as multipart field file.
curl -F "file=@sample.pdf" https://YOUR_SPACE.hf.space/ocr
Response shape: {"pages_text": ["page 1", "page 2"]}
"""
@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"