Spaces:
Sleeping
Sleeping
| 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", | |
| ) | |
| 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> | |
| """ | |
| 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(), | |
| } | |
| 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" | |