copilot-swe-agent[bot]
50573ec7be
Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
754 lines
30 KiB
Python
754 lines
30 KiB
Python
#!/usr/bin/env python3
|
||
"""OCR provider abstraction layer for DocuElevate.
|
||
|
||
This module provides a pluggable abstraction for various OCR engines, allowing
|
||
the platform to work with Azure Document Intelligence, Tesseract, EasyOCR,
|
||
Mistral OCR, Google Cloud Document AI, and AWS Textract without being locked to
|
||
a single vendor.
|
||
|
||
Provider selection is controlled by the ``OCR_PROVIDERS`` environment variable
|
||
(comma-separated list, e.g. ``azure,tesseract``). When multiple providers are
|
||
specified, all enabled providers run in parallel and the results are
|
||
cross-checked by the configured AI model to produce the best final output.
|
||
|
||
**Searchable PDF support by provider**:
|
||
|
||
+---------------------+---------------------------+-----------------------------+
|
||
| Provider | Embeds text layer in PDF? | Notes |
|
||
+=====================+===========================+=============================+
|
||
| azure | Yes | Returns PDF/A with text |
|
||
| | | layer from Document |
|
||
| | | Intelligence. |
|
||
+---------------------+---------------------------+-----------------------------+
|
||
| tesseract | No (text only) | Falls back to |
|
||
| | | ``embed_text_layer``. |
|
||
+---------------------+---------------------------+-----------------------------+
|
||
| easyocr | No (text only) | Falls back to |
|
||
| | | ``embed_text_layer``. |
|
||
+---------------------+---------------------------+-----------------------------+
|
||
| mistral | No (text only) | Falls back to |
|
||
| | | ``embed_text_layer``. |
|
||
+---------------------+---------------------------+-----------------------------+
|
||
| google_docai | No (text only) | Falls back to |
|
||
| | | ``embed_text_layer``. |
|
||
+---------------------+---------------------------+-----------------------------+
|
||
| aws_textract | No (text only) | Falls back to |
|
||
| | | ``embed_text_layer``. |
|
||
+---------------------+---------------------------+-----------------------------+
|
||
|
||
Providers that do **not** embed a text layer return ``searchable_pdf_path=None``
|
||
in their :class:`OCRResult`. The :func:`embed_text_layer` helper can be used
|
||
as a post-processing step to add a searchable text layer via ``ocrmypdf``
|
||
(which uses Tesseract for layout analysis and text positioning).
|
||
"""
|
||
|
||
import logging
|
||
import os
|
||
import shutil
|
||
import subprocess
|
||
from abc import ABC, abstractmethod
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
|
||
from app.config import settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def embed_text_layer(input_pdf_path: str, output_pdf_path: str, *, language: str = "eng") -> bool:
|
||
"""Embed a searchable text layer into a PDF using ``ocrmypdf``.
|
||
|
||
This function is used as a post-processing step for OCR providers that
|
||
return plain text only (Tesseract, EasyOCR, Mistral, Google DocAI, AWS
|
||
Textract). It calls ``ocrmypdf --skip-text`` which runs Tesseract under
|
||
the hood to detect text regions and embed an invisible text layer that
|
||
makes the PDF content selectable and searchable in PDF viewers.
|
||
|
||
Pages that already contain embedded text (e.g. from Azure Document
|
||
Intelligence) are skipped automatically by ``--skip-text``.
|
||
|
||
Args:
|
||
input_pdf_path: Absolute path to the source PDF file.
|
||
output_pdf_path: Absolute path where the searchable PDF is written.
|
||
If equal to *input_pdf_path* the file is overwritten in-place.
|
||
language: Tesseract language code(s) passed to ``ocrmypdf`` via
|
||
``-l``. Defaults to ``"eng"``. Use ``+``-separated codes for
|
||
multi-language documents, e.g. ``"eng+deu"``.
|
||
|
||
Returns:
|
||
``True`` when a searchable PDF was written successfully, ``False``
|
||
when ``ocrmypdf`` is not available or the process fails (a warning is
|
||
logged in the latter case so callers can degrade gracefully).
|
||
|
||
Raises:
|
||
FileNotFoundError: If *input_pdf_path* does not exist.
|
||
"""
|
||
if not os.path.exists(input_pdf_path):
|
||
raise FileNotFoundError(f"embed_text_layer: input file not found: {input_pdf_path}")
|
||
|
||
ocrmypdf_bin = shutil.which("ocrmypdf")
|
||
if ocrmypdf_bin is None:
|
||
logger.warning(
|
||
"ocrmypdf not found on PATH – skipping text-layer embedding. "
|
||
"Install ocrmypdf (and tesseract-ocr) to enable searchable PDF output."
|
||
)
|
||
return False
|
||
|
||
in_place = input_pdf_path == output_pdf_path
|
||
if in_place:
|
||
import tempfile
|
||
|
||
tmp_fd, tmp_path = tempfile.mkstemp(suffix=".pdf", dir=os.path.dirname(input_pdf_path))
|
||
os.close(tmp_fd)
|
||
final_output = tmp_path
|
||
else:
|
||
final_output = output_pdf_path
|
||
|
||
cmd = [
|
||
ocrmypdf_bin,
|
||
"--skip-text", # skip pages that already have a text layer (e.g. Azure output)
|
||
"--quiet", # suppress progress output; errors still appear on stderr
|
||
"-l",
|
||
language,
|
||
input_pdf_path,
|
||
final_output,
|
||
]
|
||
|
||
logger.info(f"[embed_text_layer] Running: {' '.join(cmd)}")
|
||
|
||
# Security note: shell=False (the default) is used so no shell interpolation occurs.
|
||
# ocrmypdf_bin is resolved via shutil.which() (trusted system PATH).
|
||
# input_pdf_path / final_output are internal workdir paths, not raw user input.
|
||
try:
|
||
proc = subprocess.run(cmd, capture_output=True, text=True, timeout=600, check=False) # noqa: S603
|
||
except subprocess.TimeoutExpired:
|
||
logger.warning("[embed_text_layer] ocrmypdf timed out after 600 s; skipping text-layer embedding")
|
||
if in_place and os.path.exists(final_output):
|
||
os.remove(final_output)
|
||
return False
|
||
|
||
if proc.returncode != 0:
|
||
stderr_snippet = proc.stderr.strip()[:500]
|
||
logger.warning(
|
||
f"[embed_text_layer] ocrmypdf exited with code {proc.returncode}; "
|
||
f"skipping text-layer embedding. stderr: {stderr_snippet}"
|
||
)
|
||
if in_place and os.path.exists(final_output):
|
||
os.remove(final_output)
|
||
return False
|
||
|
||
if in_place:
|
||
# Atomically replace the original file with the processed output.
|
||
os.replace(final_output, input_pdf_path)
|
||
|
||
logger.info(f"[embed_text_layer] Searchable PDF written to {output_pdf_path}")
|
||
return True
|
||
|
||
|
||
class OCRResult:
|
||
"""Container for a single OCR provider's output.
|
||
|
||
Attributes:
|
||
provider: Name of the OCR provider (e.g. ``"azure"``, ``"tesseract"``).
|
||
text: The extracted plain text.
|
||
searchable_pdf_path: Optional path to a searchable PDF produced by the
|
||
provider. ``None`` when the provider does not produce PDFs.
|
||
rotation_data: Optional dict mapping page indices to detected rotation
|
||
angles (same format used by the Azure task).
|
||
metadata: Provider-specific metadata dict (e.g. confidence scores).
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
provider: str,
|
||
text: str,
|
||
searchable_pdf_path: Optional[str] = None,
|
||
rotation_data: Optional[Dict[int, float]] = None,
|
||
metadata: Optional[Dict[str, Any]] = None,
|
||
) -> None:
|
||
self.provider = provider
|
||
self.text = text
|
||
self.searchable_pdf_path = searchable_pdf_path
|
||
self.rotation_data = rotation_data or {}
|
||
self.metadata = metadata or {}
|
||
|
||
def __repr__(self) -> str:
|
||
return (
|
||
f"OCRResult(provider={self.provider!r}, "
|
||
f"chars={len(self.text)}, "
|
||
f"has_pdf={self.searchable_pdf_path is not None})"
|
||
)
|
||
|
||
|
||
class OCRProvider(ABC):
|
||
"""Abstract base class for OCR providers.
|
||
|
||
Concrete providers must implement :meth:`process`, which accepts a path to a
|
||
PDF file and returns an :class:`OCRResult`.
|
||
|
||
Subclasses should also set :attr:`name` to a short, stable identifier
|
||
(e.g. ``"azure"``, ``"tesseract"``).
|
||
"""
|
||
|
||
#: Short, stable identifier for this provider. Must match the key used in
|
||
#: :data:`_PROVIDER_MAP` and in the ``OCR_PROVIDERS`` setting.
|
||
name: str = "unknown"
|
||
|
||
@abstractmethod
|
||
def process(self, file_path: str) -> OCRResult:
|
||
"""Run OCR on *file_path* and return an :class:`OCRResult`.
|
||
|
||
Args:
|
||
file_path: Absolute path to the input PDF file.
|
||
|
||
Returns:
|
||
An :class:`OCRResult` with the extracted text and optional
|
||
searchable-PDF path.
|
||
|
||
Raises:
|
||
Exception: If OCR processing fails.
|
||
"""
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Provider implementations
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class AzureOCRProvider(OCRProvider):
|
||
"""OCR via Azure Document Intelligence (the existing provider).
|
||
|
||
Credentials are read from ``settings.azure_ai_key`` and
|
||
``settings.azure_endpoint``.
|
||
"""
|
||
|
||
name = "azure"
|
||
|
||
def process(self, file_path: str) -> OCRResult:
|
||
from azure.ai.documentintelligence import DocumentIntelligenceClient
|
||
from azure.ai.documentintelligence.models import AnalyzeOutputOption
|
||
from azure.core.credentials import AzureKeyCredential
|
||
|
||
if not settings.azure_ai_key:
|
||
raise ValueError("AZURE_AI_KEY must be set when using the Azure OCR provider.")
|
||
if not settings.azure_endpoint:
|
||
raise ValueError("AZURE_ENDPOINT must be set when using the Azure OCR provider.")
|
||
|
||
client = DocumentIntelligenceClient(
|
||
endpoint=settings.azure_endpoint,
|
||
credential=AzureKeyCredential(settings.azure_ai_key),
|
||
)
|
||
|
||
with open(file_path, "rb") as f:
|
||
poller = client.begin_analyze_document("prebuilt-read", body=f, output=[AnalyzeOutputOption.PDF])
|
||
result = poller.result()
|
||
operation_id = poller.details["operation_id"]
|
||
|
||
# Extract rotation data
|
||
rotation_data: Dict[int, float] = {}
|
||
if hasattr(result, "pages") and result.pages:
|
||
for i, page in enumerate(result.pages):
|
||
if hasattr(page, "angle") and page.angle is not None and page.angle != 0:
|
||
rotation_data[i] = page.angle
|
||
|
||
# Retrieve searchable PDF
|
||
response = client.get_analyze_result_pdf(model_id=result.model_id, result_id=operation_id)
|
||
searchable_pdf_path = file_path # overwrite in place
|
||
with open(searchable_pdf_path, "wb") as writer:
|
||
writer.writelines(response)
|
||
|
||
extracted_text = result.content if result.content else ""
|
||
logger.info(f"[AzureOCR] Extracted {len(extracted_text)} chars from {os.path.basename(file_path)}")
|
||
|
||
return OCRResult(
|
||
provider="azure",
|
||
text=extracted_text,
|
||
searchable_pdf_path=searchable_pdf_path,
|
||
rotation_data=rotation_data,
|
||
)
|
||
|
||
|
||
class TesseractOCRProvider(OCRProvider):
|
||
"""OCR via Tesseract (self-hosted, open-source).
|
||
|
||
Requires ``pytesseract`` and ``Pillow`` to be installed, plus the
|
||
Tesseract binary on the system.
|
||
|
||
Config knobs (from :class:`~app.config.Settings`):
|
||
- ``tesseract_cmd`` – path to the ``tesseract`` binary (optional).
|
||
- ``tesseract_language`` – Tesseract language code(s), e.g. ``"eng"`` or
|
||
``"eng+deu"`` (default: ``"eng"``).
|
||
"""
|
||
|
||
name = "tesseract"
|
||
|
||
def process(self, file_path: str) -> OCRResult:
|
||
try:
|
||
import pytesseract
|
||
from pdf2image import convert_from_path
|
||
except ImportError as exc:
|
||
raise RuntimeError(
|
||
"pytesseract and pdf2image are required for the Tesseract OCR provider. "
|
||
"Install them with: pip install pytesseract pdf2image"
|
||
) from exc
|
||
|
||
tesseract_cmd = getattr(settings, "tesseract_cmd", None)
|
||
if tesseract_cmd:
|
||
pytesseract.pytesseract.tesseract_cmd = tesseract_cmd
|
||
|
||
lang = getattr(settings, "tesseract_language", None) or "eng"
|
||
|
||
logger.info(f"[TesseractOCR] Processing {os.path.basename(file_path)} (lang={lang})")
|
||
pages = convert_from_path(file_path, dpi=300)
|
||
texts: List[str] = []
|
||
for i, page_img in enumerate(pages):
|
||
page_text = pytesseract.image_to_string(page_img, lang=lang)
|
||
texts.append(page_text)
|
||
logger.debug(f"[TesseractOCR] Page {i + 1}: {len(page_text)} chars")
|
||
|
||
extracted_text = "\n".join(texts)
|
||
logger.info(f"[TesseractOCR] Extracted {len(extracted_text)} chars total")
|
||
|
||
return OCRResult(
|
||
provider="tesseract",
|
||
text=extracted_text,
|
||
)
|
||
|
||
|
||
class EasyOCRProvider(OCRProvider):
|
||
"""OCR via EasyOCR (self-hosted, deep-learning based).
|
||
|
||
Requires the ``easyocr`` package to be installed.
|
||
|
||
Config knobs (from :class:`~app.config.Settings`):
|
||
- ``easyocr_languages`` – comma-separated list of language codes
|
||
(default: ``"en"``).
|
||
- ``easyocr_gpu`` – whether to use GPU acceleration (default: ``False``).
|
||
"""
|
||
|
||
name = "easyocr"
|
||
|
||
def process(self, file_path: str) -> OCRResult:
|
||
try:
|
||
import easyocr
|
||
from pdf2image import convert_from_path
|
||
except ImportError as exc:
|
||
raise RuntimeError(
|
||
"easyocr and pdf2image are required for the EasyOCR provider. "
|
||
"Install them with: pip install easyocr pdf2image"
|
||
) from exc
|
||
|
||
lang_str = getattr(settings, "easyocr_languages", None) or "en"
|
||
langs = [lang.strip() for lang in lang_str.split(",") if lang.strip()]
|
||
gpu = getattr(settings, "easyocr_gpu", False)
|
||
|
||
logger.info(f"[EasyOCR] Processing {os.path.basename(file_path)} (langs={langs}, gpu={gpu})")
|
||
reader = easyocr.Reader(langs, gpu=gpu)
|
||
pages = convert_from_path(file_path, dpi=300)
|
||
texts: List[str] = []
|
||
for i, page_img in enumerate(pages):
|
||
result = reader.readtext(page_img, detail=0, paragraph=True)
|
||
page_text = "\n".join(result)
|
||
texts.append(page_text)
|
||
logger.debug(f"[EasyOCR] Page {i + 1}: {len(page_text)} chars")
|
||
|
||
extracted_text = "\n".join(texts)
|
||
logger.info(f"[EasyOCR] Extracted {len(extracted_text)} chars total")
|
||
|
||
return OCRResult(
|
||
provider="easyocr",
|
||
text=extracted_text,
|
||
)
|
||
|
||
|
||
class MistralOCRProvider(OCRProvider):
|
||
"""OCR via Mistral's dedicated OCR API (``/v1/ocr``).
|
||
|
||
For **PDF** files the document is first uploaded to the Mistral Files API
|
||
(``POST /v1/files``) to obtain a signed URL, then the OCR endpoint is
|
||
called with ``document_url``. For **image** files (JPEG, PNG, GIF, WEBP,
|
||
BMP, TIFF) the file is base64-encoded and passed directly as
|
||
``image_url``. Passing a PDF as a ``data:application/pdf`` data-URI to
|
||
the image path is explicitly rejected by the API and will produce a 422
|
||
error, so the two paths are kept strictly separate.
|
||
|
||
Config knobs (from :class:`~app.config.Settings`):
|
||
- ``mistral_api_key`` – Mistral API key.
|
||
- ``mistral_ocr_model`` – model name (default: ``"mistral-ocr-latest"``).
|
||
"""
|
||
|
||
name = "mistral"
|
||
|
||
# MIME types that may be sent as base64 image_url payloads
|
||
_IMAGE_MIME_TYPES: frozenset = frozenset(
|
||
{
|
||
"image/jpeg",
|
||
"image/png",
|
||
"image/gif",
|
||
"image/webp",
|
||
"image/bmp",
|
||
"image/tiff",
|
||
}
|
||
)
|
||
|
||
def process(self, file_path: str) -> OCRResult:
|
||
import base64
|
||
import mimetypes
|
||
|
||
try:
|
||
import requests as req
|
||
except ImportError as exc:
|
||
raise RuntimeError("requests package is required for the Mistral OCR provider.") from exc
|
||
|
||
api_key = getattr(settings, "mistral_api_key", None)
|
||
if not api_key:
|
||
raise ValueError("MISTRAL_API_KEY must be set when using the Mistral OCR provider.")
|
||
|
||
model = getattr(settings, "mistral_ocr_model", None) or "mistral-ocr-latest"
|
||
base_url = "https://api.mistral.ai/v1"
|
||
auth_headers: Dict[str, str] = {"Authorization": f"Bearer {api_key}"}
|
||
|
||
logger.info(f"[MistralOCR] Processing {os.path.basename(file_path)} with {model}")
|
||
|
||
# Determine MIME type from extension, falling back to magic bytes.
|
||
mime_type, _ = mimetypes.guess_type(file_path)
|
||
if mime_type is None:
|
||
with open(file_path, "rb") as fh:
|
||
magic = fh.read(5)
|
||
if magic.startswith(b"%PDF-"):
|
||
mime_type = "application/pdf"
|
||
else:
|
||
raise ValueError(
|
||
f"[MistralOCR] Cannot determine file type for '{os.path.basename(file_path)}'. "
|
||
"Supported types: PDF, JPEG, PNG, GIF, WEBP, BMP, TIFF."
|
||
)
|
||
|
||
document: Dict[str, Any]
|
||
if mime_type == "application/pdf":
|
||
document = self._upload_pdf_and_get_document(file_path, base_url, auth_headers, req)
|
||
elif mime_type in self._IMAGE_MIME_TYPES:
|
||
with open(file_path, "rb") as fh:
|
||
img_b64 = base64.b64encode(fh.read()).decode("utf-8")
|
||
document = {
|
||
"type": "image_url",
|
||
"image_url": f"data:{mime_type};base64,{img_b64}",
|
||
}
|
||
else:
|
||
raise ValueError(
|
||
f"[MistralOCR] Unsupported file type '{mime_type}' for "
|
||
f"'{os.path.basename(file_path)}'. "
|
||
"Supported types: PDF, JPEG, PNG, GIF, WEBP, BMP, TIFF."
|
||
)
|
||
|
||
ocr_payload: Dict[str, Any] = {"model": model, "document": document}
|
||
|
||
resp = req.post(
|
||
f"{base_url}/ocr",
|
||
headers={**auth_headers, "Content-Type": "application/json"},
|
||
json=ocr_payload,
|
||
timeout=300,
|
||
)
|
||
resp.raise_for_status()
|
||
ocr_data = resp.json()
|
||
|
||
pages = ocr_data.get("pages", [])
|
||
extracted_text = "\n\n".join(page.get("markdown", "") for page in pages).strip()
|
||
logger.info(f"[MistralOCR] Extracted {len(extracted_text)} chars from {len(pages)} page(s)")
|
||
|
||
return OCRResult(provider="mistral", text=extracted_text)
|
||
|
||
def _upload_pdf_and_get_document(
|
||
self,
|
||
file_path: str,
|
||
base_url: str,
|
||
auth_headers: Dict[str, str],
|
||
req: Any,
|
||
) -> Dict[str, Any]:
|
||
"""Upload *file_path* to the Mistral Files API and return an OCR document dict.
|
||
|
||
Args:
|
||
file_path: Local path to the PDF file.
|
||
base_url: Mistral API base URL.
|
||
auth_headers: Dict containing the ``Authorization`` header.
|
||
req: The ``requests`` module (injected to allow mocking in tests).
|
||
|
||
Returns:
|
||
A document dict suitable for the ``/v1/ocr`` payload, e.g.
|
||
``{"type": "document_url", "document_url": "https://..."}``.
|
||
|
||
Raises:
|
||
requests.HTTPError: If the Files API upload or URL retrieval fails.
|
||
"""
|
||
logger.info(f"[MistralOCR] Uploading '{os.path.basename(file_path)}' to Mistral Files API")
|
||
|
||
with open(file_path, "rb") as fh:
|
||
upload_resp = req.post(
|
||
f"{base_url}/files",
|
||
headers=auth_headers,
|
||
files={"file": (os.path.basename(file_path), fh, "application/pdf")},
|
||
data={"purpose": "ocr"},
|
||
timeout=300,
|
||
)
|
||
upload_resp.raise_for_status()
|
||
file_id = upload_resp.json()["id"]
|
||
|
||
logger.info(f"[MistralOCR] Uploaded file id={file_id}; fetching signed URL")
|
||
|
||
url_resp = req.get(
|
||
f"{base_url}/files/{file_id}/url",
|
||
headers=auth_headers,
|
||
params={"expiry": 24},
|
||
timeout=30,
|
||
)
|
||
url_resp.raise_for_status()
|
||
signed_url: str = url_resp.json()["url"]
|
||
|
||
return {"type": "document_url", "document_url": signed_url}
|
||
|
||
|
||
class GoogleDocAIOCRProvider(OCRProvider):
|
||
"""OCR via Google Cloud Document AI.
|
||
|
||
Config knobs (from :class:`~app.config.Settings`):
|
||
- ``google_docai_credentials_json`` – Service account JSON (optional;
|
||
falls back to ``google_drive_credentials_json`` or ADC).
|
||
- ``google_docai_project_id`` – GCP project ID (required).
|
||
- ``google_docai_processor_id`` – Document AI processor ID (required).
|
||
- ``google_docai_location`` – processor location, e.g. ``"us"`` (default:
|
||
``"us"``).
|
||
"""
|
||
|
||
name = "google_docai"
|
||
|
||
def process(self, file_path: str) -> OCRResult:
|
||
try:
|
||
from google.cloud import documentai
|
||
from google.oauth2 import service_account
|
||
except ImportError as exc:
|
||
raise RuntimeError(
|
||
"google-cloud-documentai is required for the Google Document AI OCR provider. "
|
||
"Install it with: pip install google-cloud-documentai"
|
||
) from exc
|
||
|
||
import json
|
||
|
||
project_id = getattr(settings, "google_docai_project_id", None)
|
||
processor_id = getattr(settings, "google_docai_processor_id", None)
|
||
location = getattr(settings, "google_docai_location", None) or "us"
|
||
|
||
if not project_id or not processor_id:
|
||
raise ValueError(
|
||
"GOOGLE_DOCAI_PROJECT_ID and GOOGLE_DOCAI_PROCESSOR_ID must be set "
|
||
"when using the Google Document AI OCR provider."
|
||
)
|
||
|
||
# Credentials: prefer dedicated docai key, then fall back to gdrive SA key
|
||
creds_json = getattr(settings, "google_docai_credentials_json", None) or getattr(
|
||
settings, "google_drive_credentials_json", None
|
||
)
|
||
|
||
creds = None
|
||
if creds_json:
|
||
try:
|
||
creds_info = json.loads(creds_json)
|
||
creds = service_account.Credentials.from_service_account_info(
|
||
creds_info,
|
||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"[GoogleDocAIOCR] Failed to parse credentials JSON: {e}; using ADC")
|
||
|
||
client_options = {"api_endpoint": f"{location}-documentai.googleapis.com"}
|
||
client = documentai.DocumentProcessorServiceClient(
|
||
credentials=creds,
|
||
client_options=client_options,
|
||
)
|
||
|
||
processor_name = client.processor_path(project_id, location, processor_id)
|
||
|
||
with open(file_path, "rb") as f:
|
||
raw_document = documentai.RawDocument(content=f.read(), mime_type="application/pdf")
|
||
|
||
request = documentai.ProcessRequest(name=processor_name, raw_document=raw_document)
|
||
result = client.process_document(request=request)
|
||
document = result.document
|
||
extracted_text = document.text or ""
|
||
logger.info(f"[GoogleDocAIOCR] Extracted {len(extracted_text)} chars")
|
||
|
||
return OCRResult(provider="google_docai", text=extracted_text)
|
||
|
||
|
||
class AWSTextractOCRProvider(OCRProvider):
|
||
"""OCR via AWS Textract.
|
||
|
||
Reuses existing AWS credentials from settings (``aws_access_key_id``,
|
||
``aws_secret_access_key``, ``aws_region``).
|
||
"""
|
||
|
||
name = "aws_textract"
|
||
|
||
def process(self, file_path: str) -> OCRResult:
|
||
try:
|
||
import boto3
|
||
except ImportError as exc:
|
||
raise RuntimeError("boto3 is required for the AWS Textract OCR provider.") from exc
|
||
|
||
aws_access_key_id = getattr(settings, "aws_access_key_id", None)
|
||
aws_secret_access_key = getattr(settings, "aws_secret_access_key", None)
|
||
region = getattr(settings, "aws_region", None) or "us-east-1"
|
||
|
||
if not aws_access_key_id or not aws_secret_access_key:
|
||
raise ValueError(
|
||
"AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY must be set when using the AWS Textract OCR provider."
|
||
)
|
||
|
||
logger.info(f"[AWSTextract] Processing {os.path.basename(file_path)} (region={region})")
|
||
|
||
client = boto3.client(
|
||
"textract",
|
||
aws_access_key_id=aws_access_key_id,
|
||
aws_secret_access_key=aws_secret_access_key,
|
||
region_name=region,
|
||
)
|
||
|
||
with open(file_path, "rb") as f:
|
||
document_bytes = f.read()
|
||
|
||
response = client.detect_document_text(Document={"Bytes": document_bytes})
|
||
|
||
lines: List[str] = []
|
||
for block in response.get("Blocks", []):
|
||
if block.get("BlockType") == "LINE":
|
||
text = block.get("Text", "")
|
||
if text:
|
||
lines.append(text)
|
||
|
||
extracted_text = "\n".join(lines)
|
||
logger.info(f"[AWSTextract] Extracted {len(extracted_text)} chars")
|
||
|
||
return OCRResult(provider="aws_textract", text=extracted_text)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Factory & multi-provider orchestration
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_PROVIDER_MAP: Dict[str, type] = {
|
||
"azure": AzureOCRProvider,
|
||
"tesseract": TesseractOCRProvider,
|
||
"easyocr": EasyOCRProvider,
|
||
"mistral": MistralOCRProvider,
|
||
"google_docai": GoogleDocAIOCRProvider,
|
||
"aws_textract": AWSTextractOCRProvider,
|
||
}
|
||
|
||
# Sorted list of known provider names (kept in sync with _PROVIDER_MAP)
|
||
KNOWN_OCR_PROVIDERS: List[str] = sorted(_PROVIDER_MAP.keys())
|
||
|
||
# Maximum characters per OCR result sent to the AI for merging.
|
||
# Keeping this bounded prevents excessively large prompts that would exhaust
|
||
# the model's context window or incur high token costs.
|
||
MAX_OCR_TEXT_FOR_AI_MERGE = 4000
|
||
|
||
|
||
def get_ocr_providers() -> List[OCRProvider]:
|
||
"""Return a list of configured OCR provider instances.
|
||
|
||
Reads ``settings.ocr_providers`` (comma-separated provider names) and
|
||
returns one instantiated provider per entry. Falls back to ``["azure"]``
|
||
when the setting is absent.
|
||
"""
|
||
raw = getattr(settings, "ocr_providers", None) or "azure"
|
||
provider_names = [name.strip().lower() for name in raw.split(",") if name.strip()]
|
||
|
||
providers: List[OCRProvider] = []
|
||
for name in provider_names:
|
||
cls = _PROVIDER_MAP.get(name)
|
||
if cls is None:
|
||
logger.warning(f"Unknown OCR provider '{name}' in OCR_PROVIDERS – skipping.")
|
||
continue
|
||
providers.append(cls())
|
||
logger.debug(f"Registered OCR provider: {name}")
|
||
|
||
if not providers:
|
||
logger.warning("No valid OCR providers configured, falling back to Azure.")
|
||
providers.append(AzureOCRProvider())
|
||
|
||
return providers
|
||
|
||
|
||
def merge_ocr_results(results: List[OCRResult], filename: str) -> Tuple[str, Optional[str], Dict[int, float]]:
|
||
"""Select the best text from multiple OCR results.
|
||
|
||
When only one result is available the text is returned as-is. When
|
||
multiple results exist the AI model is consulted to pick or merge the best
|
||
version (controlled by ``settings.ocr_merge_strategy``).
|
||
|
||
Args:
|
||
results: Non-empty list of :class:`OCRResult` objects.
|
||
filename: Document filename used for logging.
|
||
|
||
Returns:
|
||
A 3-tuple of ``(best_text, searchable_pdf_path, rotation_data)`` where
|
||
*searchable_pdf_path* and *rotation_data* come from the first result
|
||
that provides them.
|
||
"""
|
||
if not results:
|
||
return "", None, {}
|
||
|
||
if len(results) == 1:
|
||
r = results[0]
|
||
return r.text, r.searchable_pdf_path, r.rotation_data
|
||
|
||
strategy = getattr(settings, "ocr_merge_strategy", None) or "ai_merge"
|
||
logger.info(f"Merging {len(results)} OCR results for {filename} (strategy={strategy})")
|
||
|
||
# Best searchable PDF comes from the first provider that produced one
|
||
searchable_pdf_path = next((r.searchable_pdf_path for r in results if r.searchable_pdf_path), None)
|
||
# Best rotation data comes from the first provider that detected any
|
||
rotation_data = next((r.rotation_data for r in results if r.rotation_data), {})
|
||
|
||
if strategy == "primary":
|
||
# Simply return the first result's text
|
||
return results[0].text, searchable_pdf_path, rotation_data
|
||
|
||
if strategy == "longest":
|
||
best = max(results, key=lambda r: len(r.text))
|
||
return best.text, searchable_pdf_path, rotation_data
|
||
|
||
# Default: ai_merge – ask AI to pick/merge the best text
|
||
try:
|
||
from app.utils.ai_provider import get_ai_provider
|
||
|
||
provider = get_ai_provider()
|
||
model = settings.ai_model or settings.openai_model
|
||
|
||
extracts_block = "\n\n".join(
|
||
f"--- OCR Engine: {r.provider} ---\n{r.text[:MAX_OCR_TEXT_FOR_AI_MERGE]}" for r in results
|
||
)
|
||
messages = [
|
||
{
|
||
"role": "system",
|
||
"content": (
|
||
"You are an expert document editor. You will receive OCR extracts of the same document "
|
||
"produced by different OCR engines. Your task is to produce a single, clean, accurate "
|
||
"version of the text by cross-referencing all inputs. "
|
||
"Fix obvious OCR errors, preserve document structure, and return ONLY the final text."
|
||
),
|
||
},
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
f"Document: {filename}\n\n"
|
||
f"The following are OCR extracts from different engines:\n\n{extracts_block}\n\n"
|
||
"Please merge these into the most accurate version of the document text."
|
||
),
|
||
},
|
||
]
|
||
merged_text = provider.chat_completion(messages, model=model, temperature=0)
|
||
logger.info(f"AI-merged OCR text: {len(merged_text)} chars for {filename}")
|
||
return merged_text, searchable_pdf_path, rotation_data
|
||
except Exception as exc:
|
||
logger.error(f"AI merge failed for {filename}: {exc}; falling back to longest result")
|
||
best = max(results, key=lambda r: len(r.text))
|
||
return best.text, searchable_pdf_path, rotation_data
|