Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import io | |
| import os | |
| import re | |
| import tempfile | |
| import time | |
| import threading | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| from dataclasses import dataclass, field | |
| from pathlib import Path | |
| from typing import Iterable | |
| DIRECT_TEXT_MIN_CHARS = 30 | |
| SHORT_PDF_PAGE_LIMIT = 20 | |
| TROCR_CONFIDENCE_THRESHOLD = 0.78 | |
| DEFAULT_GEMINI_MODELS = ( | |
| "gemini-2.5-flash-lite", | |
| "gemini-2.5-flash", | |
| ) | |
| class PageResult: | |
| page_number: int | |
| engine: str | |
| text: str | |
| confidence: float | None = None | |
| class ExtractionResult: | |
| text: str | |
| route: str | |
| page_count: int | |
| direct_text_found: bool | |
| pages: list[PageResult] = field(default_factory=list) | |
| warnings: list[str] = field(default_factory=list) | |
| class PdfTextExtractor: | |
| """Implements the requested PDF-to-text decision flow.""" | |
| def __init__( | |
| self, | |
| gemini_model: str | None = None, | |
| trocr_model: str = "microsoft/trocr-base-printed", | |
| short_pdf_page_limit: int = SHORT_PDF_PAGE_LIMIT, | |
| trocr_confidence_threshold: float = TROCR_CONFIDENCE_THRESHOLD, | |
| ) -> None: | |
| self.trocr_model_name = trocr_model | |
| self.short_pdf_page_limit = short_pdf_page_limit | |
| self.trocr_confidence_threshold = trocr_confidence_threshold | |
| self._gemini_client = None | |
| self._gemini_clients = {} | |
| self._gemini_key_cursor = 0 | |
| self._gemini_key_lock = threading.Lock() | |
| self._gemini_quota_blocked_until: dict[str, float] = {} | |
| self._trocr_processor = None | |
| self._trocr_model = None | |
| self._load_local_env() | |
| self.gemini_models = self._gemini_model_candidates(gemini_model) | |
| self.gemini_max_retries = min( | |
| self._env_int("OCR_GEMINI_MAX_RETRIES", "GEMINI_MAX_RETRIES", default=1, minimum=1), | |
| 4, | |
| ) | |
| self.gemini_retry_delay = min( | |
| self._env_float("OCR_GEMINI_RETRY_DELAY", "GEMINI_RETRY_DELAY", default=0.5, minimum=0.0), | |
| 3.0, | |
| ) | |
| self.gemini_timeout_ms = int( | |
| min( | |
| self._env_float("OCR_GEMINI_TIMEOUT", "GEMINI_TIMEOUT", default=20.0, minimum=5.0), | |
| 120.0, | |
| ) | |
| * 1000 | |
| ) | |
| self.gemini_key_quota_cooldown = min( | |
| self._env_float("GEMINI_KEY_QUOTA_COOLDOWN", default=300.0, minimum=0.0), | |
| 86400.0, | |
| ) | |
| self.gemini_concurrency = min( | |
| self._env_int("OCR_GEMINI_CONCURRENCY", default=3, minimum=1), | |
| 5, | |
| ) | |
| self.render_scale = self._env_float("OCR_RENDER_SCALE", default=1.25, minimum=1.0) | |
| self.fast_ocr = os.getenv("OCR_FAST_MODE", "1") != "0" | |
| def extract(self, pdf_path: str | Path) -> ExtractionResult: | |
| path = Path(pdf_path) | |
| document = self._open_document(path) | |
| try: | |
| page_count = document.page_count | |
| direct_text = self._extract_direct_text(document) | |
| if self._has_enough_text(direct_text): | |
| return ExtractionResult( | |
| text=self._clean_text_only(direct_text), | |
| route="PyMuPDF direct text extraction", | |
| page_count=page_count, | |
| direct_text_found=True, | |
| ) | |
| images = self._convert_pages_to_preprocessed_images(document) | |
| if page_count <= self.short_pdf_page_limit: | |
| pages = [] | |
| warnings = [] | |
| for index, image_bytes in images: | |
| text, warning = self._safe_gemini_ocr(image_bytes, index) | |
| if warning: | |
| warnings.append(warning) | |
| pages.append( | |
| PageResult( | |
| page_number=index, | |
| engine="Gemini Vision OCR" if text else "Skipped non-text page", | |
| text=text, | |
| ) | |
| ) | |
| if not pages: | |
| warnings.append("No readable text pages were processed.") | |
| result = self._finalize_ocr_result( | |
| route="PyMuPDF images -> OpenCV preprocess -> Gemini Vision OCR", | |
| page_count=page_count, | |
| direct_text_found=False, | |
| pages=pages, | |
| ) | |
| result.warnings.extend(warnings) | |
| return result | |
| pages = [] | |
| warnings = [] | |
| for page_number, image_bytes in images: | |
| try: | |
| trocr_text, confidence = self._trocr_ocr(image_bytes) | |
| except Exception as exc: | |
| trocr_text = "" | |
| confidence = 0.0 | |
| warnings.append(f"Page {page_number} skipped: TrOCR failed ({exc}).") | |
| if confidence < self.trocr_confidence_threshold: | |
| gemini_text, warning = self._safe_gemini_ocr(image_bytes, page_number) | |
| if warning: | |
| warnings.append(warning) | |
| pages.append( | |
| PageResult( | |
| page_number=page_number, | |
| engine="TrOCR low confidence -> Gemini Vision OCR" if gemini_text else "Skipped non-text page", | |
| text=gemini_text, | |
| confidence=confidence, | |
| ) | |
| ) | |
| else: | |
| trocr_text = self._clean_text_only(trocr_text) | |
| pages.append( | |
| PageResult( | |
| page_number=page_number, | |
| engine="TrOCR" if trocr_text else "Skipped non-text page", | |
| text=trocr_text, | |
| confidence=confidence, | |
| ) | |
| ) | |
| if not pages: | |
| warnings.append("No pages were processed.") | |
| result = self._finalize_ocr_result( | |
| route="PyMuPDF images -> OpenCV preprocess -> TrOCR -> Gemini low-confidence fallback", | |
| page_count=page_count, | |
| direct_text_found=False, | |
| pages=pages, | |
| ) | |
| result.warnings.extend(warnings) | |
| return result | |
| finally: | |
| document.close() | |
| def stream_extract(self, pdf_path: str | Path) -> Iterable[dict]: | |
| path = Path(pdf_path) | |
| document = self._open_document(path) | |
| try: | |
| page_count = document.page_count | |
| yield {"type": "start", "page_count": page_count} | |
| short_pdf_ocr_jobs = [] | |
| for page_number, page in enumerate(document, start=1): | |
| direct_text = self._clean_text_only(page.get_text("text")) | |
| if self._has_enough_text(direct_text, min_chars=5): | |
| yield { | |
| "type": "page", | |
| "page_number": page_number, | |
| "page_count": page_count, | |
| "engine": "PyMuPDF direct text extraction", | |
| "route": "PyMuPDF direct text extraction", | |
| "direct_text_found": True, | |
| "confidence": None, | |
| "text": direct_text, | |
| } | |
| continue | |
| try: | |
| image_bytes = self._convert_page_to_preprocessed_image(page) | |
| except Exception as exc: | |
| yield self._skipped_page_event(page_number, page_count, f"Image preprocessing failed: {exc}") | |
| continue | |
| if page_count <= self.short_pdf_page_limit: | |
| short_pdf_ocr_jobs.append((page_number, image_bytes)) | |
| continue | |
| try: | |
| trocr_text, confidence = self._trocr_ocr(image_bytes) | |
| except Exception as exc: | |
| trocr_text = "" | |
| confidence = 0.0 | |
| yield self._skipped_page_event(page_number, page_count, f"TrOCR failed: {exc}", confidence) | |
| continue | |
| if confidence < self.trocr_confidence_threshold: | |
| ocr_text, warning = self._safe_gemini_ocr(image_bytes, page_number) | |
| yield { | |
| "type": "page", | |
| "page_number": page_number, | |
| "page_count": page_count, | |
| "engine": "TrOCR low confidence -> Gemini Vision OCR" if ocr_text else "Skipped non-text page", | |
| "route": "PyMuPDF image -> OpenCV preprocess -> TrOCR -> Gemini fallback", | |
| "direct_text_found": False, | |
| "confidence": confidence, | |
| "text": ocr_text, | |
| "warning": warning, | |
| } | |
| else: | |
| trocr_text = self._clean_text_only(trocr_text) | |
| if not self._has_enough_text(trocr_text, min_chars=2): | |
| yield self._skipped_page_event(page_number, page_count, "No readable text found.", confidence) | |
| continue | |
| yield { | |
| "type": "page", | |
| "page_number": page_number, | |
| "page_count": page_count, | |
| "engine": "TrOCR", | |
| "route": "PyMuPDF image -> OpenCV preprocess -> TrOCR", | |
| "direct_text_found": False, | |
| "confidence": confidence, | |
| "text": self._clean_text(trocr_text), | |
| } | |
| if short_pdf_ocr_jobs: | |
| yield from self._stream_gemini_ocr_jobs(short_pdf_ocr_jobs, page_count) | |
| yield {"type": "done", "page_count": page_count} | |
| finally: | |
| document.close() | |
| def _open_document(self, pdf_path: Path): | |
| try: | |
| import fitz | |
| except ImportError as exc: | |
| raise RuntimeError("PyMuPDF is required. Install it with: pip install pymupdf") from exc | |
| try: | |
| document = fitz.open(pdf_path) | |
| if document.is_encrypted: | |
| raise RuntimeError("This PDF is password-protected. Please upload an unlocked PDF.") | |
| if document.page_count == 0: | |
| raise RuntimeError("This PDF has no pages.") | |
| return document | |
| except RuntimeError: | |
| raise | |
| except Exception as exc: | |
| raise RuntimeError(f"Could not open PDF. It may be corrupted or unsupported: {exc}") from exc | |
| def _extract_direct_text(self, document) -> str: | |
| chunks = [] | |
| for page in document: | |
| chunks.append(page.get_text("text")) | |
| return "\n\n".join(chunks) | |
| def _has_enough_text(self, text: str, min_chars: int = DIRECT_TEXT_MIN_CHARS) -> bool: | |
| normalized = re.sub(r"\s+", "", text or "") | |
| return len(normalized) >= min_chars | |
| def _convert_pages_to_preprocessed_images(self, document) -> list[tuple[int, bytes]]: | |
| images = [] | |
| for index, page in enumerate(document, start=1): | |
| try: | |
| images.append((index, self._convert_page_to_preprocessed_image(page))) | |
| except Exception: | |
| continue | |
| return images | |
| def _convert_page_to_preprocessed_image(self, page) -> bytes: | |
| try: | |
| import cv2 | |
| import fitz | |
| import numpy as np | |
| from PIL import Image | |
| except ImportError as exc: | |
| raise RuntimeError( | |
| "OCR image preprocessing needs opencv-python, numpy, Pillow, and PyMuPDF." | |
| ) from exc | |
| matrix = fitz.Matrix(self.render_scale, self.render_scale) | |
| pixmap = page.get_pixmap(matrix=matrix, alpha=False) | |
| pil_image = Image.open(io.BytesIO(pixmap.tobytes("png"))).convert("RGB") | |
| array = np.array(pil_image) | |
| gray = cv2.cvtColor(array, cv2.COLOR_RGB2GRAY) | |
| if self.fast_ocr: | |
| _, thresholded = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) | |
| else: | |
| denoised = cv2.fastNlMeansDenoising(gray, h=10) | |
| thresholded = cv2.adaptiveThreshold( | |
| denoised, | |
| 255, | |
| cv2.ADAPTIVE_THRESH_GAUSSIAN_C, | |
| cv2.THRESH_BINARY, | |
| 31, | |
| 11, | |
| ) | |
| output = Image.fromarray(thresholded).convert("RGB") | |
| buffer = io.BytesIO() | |
| output.save(buffer, format="JPEG", quality=72, optimize=True) | |
| return buffer.getvalue() | |
| def _gemini_ocr(self, image_bytes: bytes) -> str: | |
| try: | |
| from google import genai | |
| from google.genai import types | |
| except ImportError as exc: | |
| raise RuntimeError("Gemini OCR requires: pip install google-genai") from exc | |
| api_keys = self._gemini_api_keys() | |
| if not api_keys: | |
| raise RuntimeError("Set GEMINI_API_KEY before using Gemini Vision OCR.") | |
| prompt = ( | |
| "Extract all readable text from this preprocessed PDF page image. " | |
| "Preserve natural reading order, headings, bullet points, tables as plain text, " | |
| "and do not add commentary. Ignore photos, diagrams, icons, borders, handwriting-like noise, " | |
| "and decorative/non-text visual content. Do not describe images. If there is no readable text, " | |
| "return an empty response." | |
| ) | |
| last_error: Exception | None = None | |
| attempted_keys = 0 | |
| quota_blocked_keys = 0 | |
| for api_key in self._gemini_key_attempt_order(api_keys): | |
| attempted_keys += 1 | |
| client = self._gemini_clients.get(api_key) | |
| if client is None: | |
| client = genai.Client( | |
| api_key=api_key, | |
| http_options=types.HttpOptions(timeout=self.gemini_timeout_ms), | |
| ) | |
| self._gemini_clients[api_key] = client | |
| quota_hit = False | |
| for model in self.gemini_models: | |
| for attempt in range(1, self.gemini_max_retries + 1): | |
| try: | |
| response = client.models.generate_content( | |
| model=model, | |
| contents=[ | |
| types.Part.from_bytes(data=image_bytes, mime_type="image/jpeg"), | |
| prompt, | |
| ], | |
| ) | |
| self._advance_gemini_key_cursor(api_keys, api_key) | |
| return (getattr(response, "text", None) or "").strip() | |
| except Exception as exc: | |
| last_error = exc | |
| if self._is_quota_error(exc): | |
| self._mark_gemini_key_quota_exhausted(api_keys, api_key) | |
| quota_blocked_keys += 1 | |
| quota_hit = True | |
| break | |
| if not self._is_retryable_gemini_error(exc): | |
| break | |
| if attempt < self.gemini_max_retries: | |
| time.sleep(self.gemini_retry_delay * attempt) | |
| if quota_hit: | |
| break | |
| models = ", ".join(self.gemini_models) | |
| raise RuntimeError( | |
| "Gemini OCR is temporarily unavailable after retries. " | |
| f"Tried {attempted_keys} active API key(s) out of {len(api_keys)} configured key(s); " | |
| f"{quota_blocked_keys} key(s) returned quota/rate-limit errors. Models tried: {models}. " | |
| f"Last error: {last_error}" | |
| ) | |
| def _safe_gemini_ocr(self, image_bytes: bytes, page_number: int) -> tuple[str, str | None]: | |
| try: | |
| text = self._clean_text_only(self._gemini_ocr(image_bytes)) | |
| except Exception as exc: | |
| return "", f"Page {page_number} skipped: OCR failed ({exc})." | |
| if not self._has_enough_text(text, min_chars=2): | |
| return "", f"Page {page_number} skipped: no readable text found." | |
| return text, None | |
| def _skipped_page_event( | |
| self, | |
| page_number: int, | |
| page_count: int, | |
| warning: str, | |
| confidence: float | None = None, | |
| ) -> dict: | |
| return { | |
| "type": "page", | |
| "page_number": page_number, | |
| "page_count": page_count, | |
| "engine": "Skipped non-text page", | |
| "route": "Skipped non-text content", | |
| "direct_text_found": False, | |
| "confidence": confidence, | |
| "text": "", | |
| "warning": warning, | |
| } | |
| def _stream_gemini_ocr_jobs(self, jobs: list[tuple[int, bytes]], page_count: int) -> Iterable[dict]: | |
| worker_count = min(self.gemini_concurrency, len(jobs)) | |
| if worker_count <= 1: | |
| for page_number, image_bytes in jobs: | |
| yield self._gemini_page_event(page_number, page_count, image_bytes) | |
| return | |
| with ThreadPoolExecutor(max_workers=worker_count) as executor: | |
| futures = { | |
| executor.submit(self._gemini_page_event, page_number, page_count, image_bytes): page_number | |
| for page_number, image_bytes in jobs | |
| } | |
| for future in as_completed(futures): | |
| try: | |
| yield future.result() | |
| except Exception as exc: | |
| yield self._skipped_page_event(futures[future], page_count, f"OCR failed: {exc}") | |
| def _gemini_page_event(self, page_number: int, page_count: int, image_bytes: bytes) -> dict: | |
| ocr_text, warning = self._safe_gemini_ocr(image_bytes, page_number) | |
| return { | |
| "type": "page", | |
| "page_number": page_number, | |
| "page_count": page_count, | |
| "engine": "Gemini Vision OCR" if ocr_text else "Skipped non-text page", | |
| "route": "PyMuPDF image -> OpenCV preprocess -> Gemini Vision OCR", | |
| "direct_text_found": False, | |
| "confidence": None, | |
| "text": ocr_text, | |
| "warning": warning, | |
| } | |
| def _gemini_api_keys(self) -> list[str]: | |
| keys = [] | |
| for key in sorted(os.environ, key=self._gemini_env_key_sort): | |
| value = os.environ[key] | |
| if self._is_gemini_key_name(key): | |
| self._add_gemini_key_values(keys, value) | |
| for file_name in (".env", "env"): | |
| env_path = Path(file_name) | |
| if not env_path.exists(): | |
| continue | |
| for line in env_path.read_text(encoding="utf-8", errors="ignore").splitlines(): | |
| stripped = line.strip() | |
| if not stripped or stripped.startswith("#") or "=" not in stripped: | |
| continue | |
| key, value = stripped.split("=", 1) | |
| key = key.strip() | |
| value = value.strip().strip('"').strip("'") | |
| if self._is_gemini_key_name(key): | |
| self._add_gemini_key_values(keys, value) | |
| return keys | |
| def _gemini_key_attempt_order(self, api_keys: list[str]) -> list[str]: | |
| with self._gemini_key_lock: | |
| now = time.time() | |
| active_keys = [ | |
| api_key | |
| for api_key in api_keys | |
| if self._gemini_quota_blocked_until.get(api_key, 0.0) <= now | |
| ] | |
| if not active_keys: | |
| active_keys = api_keys | |
| active_key_set = set(active_keys) | |
| start = self._gemini_key_cursor % len(api_keys) | |
| ordered_keys = api_keys[start:] + api_keys[:start] | |
| return [api_key for api_key in ordered_keys if api_key in active_key_set] | |
| def _advance_gemini_key_cursor(self, api_keys: list[str], api_key: str) -> None: | |
| with self._gemini_key_lock: | |
| try: | |
| self._gemini_key_cursor = (api_keys.index(api_key) + 1) % len(api_keys) | |
| except ValueError: | |
| self._gemini_key_cursor = 0 | |
| def _mark_gemini_key_quota_exhausted(self, api_keys: list[str], api_key: str) -> None: | |
| with self._gemini_key_lock: | |
| self._gemini_quota_blocked_until[api_key] = time.time() + self.gemini_key_quota_cooldown | |
| try: | |
| self._gemini_key_cursor = (api_keys.index(api_key) + 1) % len(api_keys) | |
| except ValueError: | |
| self._gemini_key_cursor = 0 | |
| def _is_gemini_key_name(self, key: str) -> bool: | |
| return ( | |
| key == "GEMINI_API_KEY" | |
| or key == "GOOGLE_API_KEY" | |
| or key == "GEMINI_API_KEYS" | |
| or key.startswith("GEMINI_API_KEY_") | |
| ) | |
| def _gemini_env_key_sort(self, key: str) -> tuple[int, int, str]: | |
| match = re.fullmatch(r"GEMINI_API_KEY_(\d+)", key) | |
| if match: | |
| return (0, int(match.group(1)), key) | |
| if key == "GEMINI_API_KEY": | |
| return (1, 0, key) | |
| if key == "GEMINI_API_KEYS": | |
| return (2, 0, key) | |
| if key == "GOOGLE_API_KEY": | |
| return (3, 0, key) | |
| return (4, 0, key) | |
| def _add_gemini_key_values(self, keys: list[str], value: str) -> None: | |
| for api_key in (part.strip() for part in value.split(",")): | |
| if api_key and api_key not in keys: | |
| keys.append(api_key) | |
| def _gemini_model_candidates(self, configured_model: str | None) -> list[str]: | |
| configured_models = os.getenv("GEMINI_MODELS") or configured_model | |
| if configured_models: | |
| models = [ | |
| model.strip() | |
| for model in configured_models.split(",") | |
| if model.strip() | |
| ] | |
| if models: | |
| return models | |
| return list(DEFAULT_GEMINI_MODELS) | |
| def _is_retryable_gemini_error(self, error: Exception) -> bool: | |
| message = str(error).lower() | |
| retryable_markers = ( | |
| "503", | |
| "unavailable", | |
| "high demand", | |
| "429", | |
| "resource_exhausted", | |
| "rate limit", | |
| "quota", | |
| "deadline", | |
| "timeout", | |
| ) | |
| return any(marker in message for marker in retryable_markers) | |
| def _is_quota_error(self, error: Exception | None) -> bool: | |
| if error is None: | |
| return False | |
| message = str(error).lower() | |
| return "429" in message or "resource_exhausted" in message or "quota" in message | |
| def _load_local_env(self) -> None: | |
| for file_name in (".env", "env"): | |
| env_path = Path(file_name) | |
| if not env_path.exists(): | |
| continue | |
| for line in env_path.read_text(encoding="utf-8", errors="ignore").splitlines(): | |
| stripped = line.strip() | |
| if not stripped or stripped.startswith("#") or "=" not in stripped: | |
| continue | |
| key, value = stripped.split("=", 1) | |
| key = key.strip() | |
| value = value.strip().strip('"').strip("'") | |
| if key and key not in os.environ: | |
| os.environ[key] = value | |
| def _env_int(self, key: str, fallback_key: str | None = None, default: int = 0, minimum: int = 0) -> int: | |
| raw_value = os.getenv(key) | |
| if raw_value is None and fallback_key: | |
| raw_value = os.getenv(fallback_key) | |
| if raw_value is None: | |
| return default | |
| try: | |
| value = int(raw_value) | |
| except ValueError: | |
| return default | |
| return max(value, minimum) | |
| def _env_float(self, key: str, fallback_key: str | None = None, default: float = 0.0, minimum: float = 0.0) -> float: | |
| raw_value = os.getenv(key) | |
| if raw_value is None and fallback_key: | |
| raw_value = os.getenv(fallback_key) | |
| if raw_value is None: | |
| return default | |
| try: | |
| value = float(raw_value) | |
| except ValueError: | |
| return default | |
| return max(value, minimum) | |
| def _trocr_ocr(self, image_bytes: bytes) -> tuple[str, float]: | |
| try: | |
| import torch | |
| from PIL import Image | |
| from transformers import TrOCRProcessor, VisionEncoderDecoderModel | |
| except ImportError as exc: | |
| raise RuntimeError( | |
| "TrOCR requires: pip install torch transformers Pillow" | |
| ) from exc | |
| if self._trocr_processor is None or self._trocr_model is None: | |
| self._trocr_processor = TrOCRProcessor.from_pretrained(self.trocr_model_name) | |
| self._trocr_model = VisionEncoderDecoderModel.from_pretrained(self.trocr_model_name) | |
| image = Image.open(io.BytesIO(image_bytes)).convert("RGB") | |
| pixel_values = self._trocr_processor(images=image, return_tensors="pt").pixel_values | |
| with torch.no_grad(): | |
| generated = self._trocr_model.generate( | |
| pixel_values, | |
| max_new_tokens=512, | |
| output_scores=True, | |
| return_dict_in_generate=True, | |
| ) | |
| text = self._trocr_processor.batch_decode( | |
| generated.sequences, | |
| skip_special_tokens=True, | |
| )[0].strip() | |
| confidence = self._score_trocr_confidence(generated.scores) | |
| return text, confidence | |
| def _score_trocr_confidence(self, scores: Iterable) -> float: | |
| try: | |
| import torch | |
| except ImportError: | |
| return 0.0 | |
| confidences = [] | |
| for score in scores: | |
| probabilities = torch.softmax(score, dim=-1) | |
| confidences.append(float(probabilities.max())) | |
| if not confidences: | |
| return 0.0 | |
| return sum(confidences) / len(confidences) | |
| def _finalize_ocr_result( | |
| self, | |
| route: str, | |
| page_count: int, | |
| direct_text_found: bool, | |
| pages: list[PageResult], | |
| ) -> ExtractionResult: | |
| merged_text = "\n\n".join( | |
| f"--- Page {page.page_number} ---\n{page.text.strip()}" | |
| for page in pages | |
| if page.text.strip() | |
| ) | |
| return ExtractionResult( | |
| text=self._clean_text(merged_text), | |
| route=route, | |
| page_count=page_count, | |
| direct_text_found=direct_text_found, | |
| pages=pages, | |
| ) | |
| def _clean_text(self, text: str) -> str: | |
| text = text.replace("\x00", "") | |
| text = re.sub(r"[ \t]+", " ", text) | |
| text = re.sub(r"\n{3,}", "\n\n", text) | |
| text = "\n".join(line.rstrip() for line in text.splitlines()) | |
| return text.strip() | |
| def _clean_text_only(self, text: str) -> str: | |
| cleaned = self._clean_text(text) | |
| lines = [] | |
| for line in cleaned.splitlines(): | |
| stripped = line.strip() | |
| if not stripped: | |
| lines.append("") | |
| continue | |
| if self._looks_like_visual_description(stripped): | |
| continue | |
| lines.append(stripped) | |
| text_only = self._clean_text("\n".join(lines)) | |
| return text_only if self._has_real_text_signal(text_only) else "" | |
| def _looks_like_visual_description(self, line: str) -> bool: | |
| normalized = re.sub(r"\s+", " ", line.lower()).strip() | |
| visual_description_patterns = ( | |
| r"^(the|this|an|a)\s+(image|photo|picture|diagram|figure|chart|graph|illustration|logo|icon)\s+", | |
| r"^(the|this)\s+page\s+(contains|shows|appears|has|is)\s+", | |
| r"^(it|this)\s+(shows|appears|looks like|contains|depicts)\s+", | |
| r"^(i can see|there is|there are)\s+", | |
| r"\b(no readable text|no text|cannot extract|not able to extract)\b", | |
| r"\b(image shows|picture shows|diagram shows|photo shows|chart shows|graph shows)\b", | |
| ) | |
| if any(re.search(pattern, normalized) for pattern in visual_description_patterns): | |
| return True | |
| visual_words = { | |
| "image", | |
| "photo", | |
| "picture", | |
| "diagram", | |
| "illustration", | |
| "visual", | |
| "icon", | |
| "logo", | |
| "background", | |
| "border", | |
| "shape", | |
| "graphic", | |
| } | |
| tokens = re.findall(r"[a-z0-9]+", normalized) | |
| if not tokens: | |
| return True | |
| visual_count = sum(1 for token in tokens if token in visual_words) | |
| return visual_count >= 2 and len(tokens) <= 18 | |
| def _has_real_text_signal(self, text: str) -> bool: | |
| normalized = re.sub(r"\s+", " ", text or "").strip() | |
| if not normalized: | |
| return False | |
| alnum_count = sum(char.isalnum() for char in normalized) | |
| alpha_count = sum(char.isalpha() for char in normalized) | |
| return alnum_count >= 2 and alpha_count >= 2 | |
| def extract_pdf_to_text(pdf_path: str | Path) -> ExtractionResult: | |
| return PdfTextExtractor().extract(pdf_path) | |
| def extract_uploaded_pdf(file_name: str, file_bytes: bytes) -> ExtractionResult: | |
| suffix = Path(file_name).suffix or ".pdf" | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file: | |
| temp_file.write(file_bytes) | |
| temp_path = Path(temp_file.name) | |
| try: | |
| return extract_pdf_to_text(temp_path) | |
| finally: | |
| temp_path.unlink(missing_ok=True) | |
| def stream_uploaded_pdf(file_name: str, file_bytes: bytes) -> Iterable[dict]: | |
| suffix = Path(file_name).suffix or ".pdf" | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file: | |
| temp_file.write(file_bytes) | |
| temp_path = Path(temp_file.name) | |
| try: | |
| yield from PdfTextExtractor().stream_extract(temp_path) | |
| finally: | |
| temp_path.unlink(missing_ok=True) | |