SmartNotes_AI / pdf_to_text.py
Urmila Saini
add DockerFile
b9fa4a6
Raw
History Blame Contribute Delete
29.9 kB
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",
)
@dataclass
class PageResult:
page_number: int
engine: str
text: str
confidence: float | None = None
@dataclass
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)