Merge pull request #390 from christianlouis/copilot/embed-ocr-text-layer-pdf
feat(ocr): embed searchable text layer for providers without native PDF output
This commit is contained in:
+12
@@ -16,6 +16,18 @@ WORKDIR /app
|
||||
COPY --from=builder /usr/local/lib/python3.14/site-packages /usr/local/lib/python3.14/site-packages
|
||||
COPY --from=builder /usr/local/bin /usr/local/bin
|
||||
|
||||
# Install system-level OCR tools required for local OCR workflows:
|
||||
# tesseract-ocr – OCR engine used by pytesseract and ocrmypdf
|
||||
# ghostscript – required by ocrmypdf for PDF/PS operations
|
||||
# poppler-utils – provides pdfinfo/pdftoppm used by pdf2image
|
||||
# unpaper – optional deskewing pre-processor used by ocrmypdf
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
tesseract-ocr \
|
||||
ghostscript \
|
||||
poppler-utils \
|
||||
unpaper \
|
||||
&& apt-get clean && rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Copy application code
|
||||
COPY ./app /app/app
|
||||
COPY ./frontend /app/frontend
|
||||
|
||||
@@ -13,6 +13,18 @@ WORKDIR /app
|
||||
COPY --from=builder /usr/local/lib/python3.14/site-packages /usr/local/lib/python3.14/site-packages
|
||||
COPY --from=builder /usr/local/bin /usr/local/bin
|
||||
|
||||
# Install system-level OCR tools required for local OCR workflows:
|
||||
# tesseract-ocr – OCR engine used by pytesseract and ocrmypdf
|
||||
# ghostscript – required by ocrmypdf for PDF/PS operations
|
||||
# poppler-utils – provides pdfinfo/pdftoppm used by pdf2image
|
||||
# unpaper – optional deskewing pre-processor used by ocrmypdf
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
tesseract-ocr \
|
||||
ghostscript \
|
||||
poppler-utils \
|
||||
unpaper \
|
||||
&& apt-get clean && rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY ./app /app/app
|
||||
COPY ./frontend /app/frontend
|
||||
COPY ./LICENSE /app/LICENSE
|
||||
|
||||
@@ -21,7 +21,7 @@ from app.config import settings
|
||||
from app.tasks.retry_config import BaseTaskWithRetry
|
||||
from app.tasks.rotate_pdf_pages import rotate_pdf_pages
|
||||
from app.utils import log_task_progress
|
||||
from app.utils.ocr_provider import OCRResult, get_ocr_providers, merge_ocr_results
|
||||
from app.utils.ocr_provider import OCRResult, embed_text_layer, get_ocr_providers, merge_ocr_results
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -107,6 +107,37 @@ def process_with_ocr(self, filename: str, file_id: Optional[int] = None):
|
||||
f"rotations={len(rotation_data)}"
|
||||
)
|
||||
|
||||
# If no provider produced a searchable PDF, post-process the original
|
||||
# PDF with ocrmypdf to embed an invisible text layer so the output is
|
||||
# selectable/searchable in PDF viewers.
|
||||
if searchable_pdf_path is None:
|
||||
lang = getattr(settings, "tesseract_language", None) or "eng"
|
||||
log_task_progress(
|
||||
task_id,
|
||||
"embed_text_layer",
|
||||
"in_progress",
|
||||
"Embedding searchable text layer into PDF",
|
||||
file_id=file_id,
|
||||
)
|
||||
embedded = embed_text_layer(tmp_file_path, tmp_file_path, language=lang)
|
||||
if embedded:
|
||||
searchable_pdf_path = tmp_file_path
|
||||
log_task_progress(
|
||||
task_id,
|
||||
"embed_text_layer",
|
||||
"success",
|
||||
"Searchable text layer embedded via ocrmypdf",
|
||||
file_id=file_id,
|
||||
)
|
||||
else:
|
||||
log_task_progress(
|
||||
task_id,
|
||||
"embed_text_layer",
|
||||
"skipped",
|
||||
"ocrmypdf unavailable – PDF will not have a searchable text layer",
|
||||
file_id=file_id,
|
||||
)
|
||||
|
||||
log_task_progress(
|
||||
task_id,
|
||||
"process_with_ocr",
|
||||
|
||||
@@ -10,10 +10,42 @@ 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
|
||||
|
||||
@@ -22,6 +54,96 @@ 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.
|
||||
|
||||
|
||||
+100
-1
@@ -446,7 +446,106 @@ OPENAI_API_KEY=sk-ant-... # passed as the api_key to LiteLLM
|
||||
|
||||
---
|
||||
|
||||
### Azure Document Intelligence
|
||||
### OCR Providers
|
||||
|
||||
DocuElevate supports multiple OCR engines that can be used individually or in combination. Configure the list of active providers with `OCR_PROVIDERS` and tune each provider with the settings below.
|
||||
|
||||
#### Provider Selection
|
||||
|
||||
| **Variable** | **Description** | **Default** |
|
||||
|-----------------------|---------------------------------------------------------------------------------------------------|-------------|
|
||||
| `OCR_PROVIDERS` | Comma-separated list of OCR engines to use, e.g. `azure`, `mistral`, `azure,tesseract`. | `azure` |
|
||||
| `OCR_MERGE_STRATEGY` | Strategy for combining results from multiple providers: `ai_merge`, `longest`, or `primary`. | `ai_merge` |
|
||||
|
||||
**Supported `OCR_PROVIDERS` values**: `azure`, `tesseract`, `easyocr`, `mistral`, `google_docai`, `aws_textract`
|
||||
|
||||
When multiple providers are listed, all run in parallel and their results are merged according to `OCR_MERGE_STRATEGY`.
|
||||
|
||||
#### Searchable PDF Text Layer
|
||||
|
||||
Not all OCR providers embed a searchable text layer in the output PDF. The table below summarises each provider's behaviour and how DocuElevate handles it:
|
||||
|
||||
| **Provider** | **Embeds text layer?** | **Notes** |
|
||||
|-------------------|------------------------|-----------|
|
||||
| `azure` | ✅ Yes | Azure Document Intelligence returns a PDF/A with an embedded text layer. |
|
||||
| `tesseract` | ❌ No (text only) | Text is extracted but the PDF is not modified. `embed_text_layer` post-processing is applied automatically. |
|
||||
| `easyocr` | ❌ No (text only) | Same as above. |
|
||||
| `mistral` | ❌ No (text only) | Mistral OCR API returns plain text; `embed_text_layer` post-processing is applied automatically. |
|
||||
| `google_docai` | ❌ No (text only) | Google Cloud Document AI returns plain text; `embed_text_layer` post-processing is applied automatically. |
|
||||
| `aws_textract` | ❌ No (text only) | AWS Textract returns plain text; `embed_text_layer` post-processing is applied automatically. |
|
||||
|
||||
For providers that do **not** embed a text layer, DocuElevate automatically runs `ocrmypdf --skip-text` after OCR to add an invisible Tesseract-generated text layer to the PDF. This makes the file selectable and searchable in PDF viewers. The step is silently skipped if `ocrmypdf` is not available on `PATH` (a warning is logged).
|
||||
|
||||
#### Azure Document Intelligence
|
||||
|
||||
| **Variable** | **Description** | **How to Obtain** |
|
||||
|-------------------------------------------|----------------------------------------------------------|-----------------------------------------|
|
||||
| `AZURE_DOCUMENT_INTELLIGENCE_KEY` | Azure Document Intelligence API key for OCR. | [Azure Portal](https://portal.azure.com/) |
|
||||
| `AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT` | Endpoint URL for Azure Document Intelligence API. | [Azure Portal](https://portal.azure.com/) |
|
||||
|
||||
#### Tesseract (self-hosted)
|
||||
|
||||
Requires `tesseract-ocr` to be installed in the Docker image or on the host. The default Docker image ships with Tesseract.
|
||||
|
||||
| **Variable** | **Description** | **Default** |
|
||||
|------------------------|-----------------------------------------------------------------------------------|-------------|
|
||||
| `TESSERACT_CMD` | Path to the `tesseract` binary (optional; auto-detected from `PATH`). | *(auto)* |
|
||||
| `TESSERACT_LANGUAGE` | Tesseract language code(s), e.g. `eng`, `eng+deu`, `deu`. | `eng+deu` |
|
||||
|
||||
```bash
|
||||
OCR_PROVIDERS=tesseract
|
||||
TESSERACT_LANGUAGE=eng+deu
|
||||
```
|
||||
|
||||
#### EasyOCR (self-hosted)
|
||||
|
||||
Requires the `easyocr` Python package. Install it separately as it is not included in the base requirements.
|
||||
|
||||
| **Variable** | **Description** | **Default** |
|
||||
|-----------------------|------------------------------------------------------------------------|-------------|
|
||||
| `EASYOCR_LANGUAGES` | Comma-separated EasyOCR language codes, e.g. `en,de,fr`. | `en,de` |
|
||||
| `EASYOCR_GPU` | Enable GPU acceleration for EasyOCR (`true`/`false`). | `false` |
|
||||
|
||||
#### Mistral OCR
|
||||
|
||||
| **Variable** | **Description** | **How to Obtain** |
|
||||
|------------------------|------------------------------------------------|------------------------------------------------|
|
||||
| `MISTRAL_API_KEY` | Mistral API key. | [console.mistral.ai](https://console.mistral.ai) |
|
||||
| `MISTRAL_OCR_MODEL` | Mistral OCR model name. | `mistral-ocr-latest` |
|
||||
|
||||
#### Google Cloud Document AI
|
||||
|
||||
| **Variable** | **Description** | **Default** |
|
||||
|----------------------------------|---------------------------------------------------------------------------------------|-------------|
|
||||
| `GOOGLE_DOCAI_PROJECT_ID` | GCP project ID (required). | *(required)* |
|
||||
| `GOOGLE_DOCAI_PROCESSOR_ID` | Document AI processor ID (required). | *(required)* |
|
||||
| `GOOGLE_DOCAI_LOCATION` | Processor location, e.g. `us` or `eu`. | `us` |
|
||||
| `GOOGLE_DOCAI_CREDENTIALS_JSON` | Service account JSON (optional; falls back to `GOOGLE_DRIVE_CREDENTIALS_JSON`). | *(optional)* |
|
||||
|
||||
#### AWS Textract
|
||||
|
||||
Reuses the AWS credentials already configured for S3 integration.
|
||||
|
||||
| **Variable** | **Description** |
|
||||
|---------------------------|--------------------------------------|
|
||||
| `AWS_ACCESS_KEY_ID` | AWS access key ID. |
|
||||
| `AWS_SECRET_ACCESS_KEY` | AWS secret access key. |
|
||||
| `AWS_REGION` | AWS region, e.g. `us-east-1`. |
|
||||
|
||||
#### Multi-Provider Example
|
||||
|
||||
```bash
|
||||
# Use both Azure (for accuracy) and Tesseract (for redundancy); merge via AI
|
||||
OCR_PROVIDERS=azure,tesseract
|
||||
OCR_MERGE_STRATEGY=ai_merge
|
||||
AZURE_AI_KEY=...
|
||||
AZURE_ENDPOINT=https://...
|
||||
TESSERACT_LANGUAGE=eng+deu
|
||||
```
|
||||
|
||||
### Azure Document Intelligence (Legacy)
|
||||
|
||||
> **Note:** This section documents the standalone Azure Document Intelligence credentials. When using `OCR_PROVIDERS=azure` these same credentials are used automatically.
|
||||
|
||||
| **Variable** | **Description** | **How to Obtain** |
|
||||
|---------------------------------|------------------------------------------|--------------------------------------------------------------------------|
|
||||
|
||||
+2
-1
@@ -41,4 +41,5 @@ litellm>=1.0.0,<2.0.0
|
||||
|
||||
# Self-hosted OCR engines (optional – only required when the provider is enabled)
|
||||
pytesseract>=0.3.10 # Python wrapper for Tesseract OCR
|
||||
pdf2image>=1.17.0 # Convert PDF pages to images (used by Tesseract and EasyOCR providers)
|
||||
pdf2image>=1.17.0 # Convert PDF pages to images (used by Tesseract and EasyOCR providers)
|
||||
ocrmypdf>=16.0.0,<17.0.0 # Post-processing: embeds searchable text layers into PDFs via Tesseract
|
||||
|
||||
@@ -1020,3 +1020,287 @@ class TestMistralOCRProvider:
|
||||
|
||||
with pytest.raises(HTTPError):
|
||||
provider.process(str(pdf_file))
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestEmbedTextLayer:
|
||||
"""Tests for the embed_text_layer utility function."""
|
||||
|
||||
def _make_pdf(self, tmp_path, name: str = "test.pdf") -> str:
|
||||
"""Create a minimal but valid-ish PDF file for testing."""
|
||||
pdf_path = tmp_path / name
|
||||
pdf_path.write_bytes(
|
||||
b"%PDF-1.4\n"
|
||||
b"1 0 obj\n<</Type /Catalog /Pages 2 0 R>>\nendobj\n"
|
||||
b"2 0 obj\n<</Type /Pages /Kids [3 0 R] /Count 1>>\nendobj\n"
|
||||
b"3 0 obj\n<</Type /Page /Parent 2 0 R /MediaBox [0 0 612 792]>>\nendobj\n"
|
||||
b"xref\n0 4\n"
|
||||
b"0000000000 65535 f \n"
|
||||
b"0000000009 00000 n \n"
|
||||
b"0000000058 00000 n \n"
|
||||
b"0000000115 00000 n \n"
|
||||
b"trailer\n<</Size 4 /Root 1 0 R>>\n"
|
||||
b"startxref\n190\n%%EOF\n"
|
||||
)
|
||||
return str(pdf_path)
|
||||
|
||||
def test_missing_input_raises_file_not_found(self, tmp_path):
|
||||
"""FileNotFoundError is raised when the input PDF does not exist."""
|
||||
from app.utils.ocr_provider import embed_text_layer
|
||||
|
||||
with pytest.raises(FileNotFoundError, match="embed_text_layer"):
|
||||
embed_text_layer("/nonexistent/path.pdf", str(tmp_path / "out.pdf"))
|
||||
|
||||
def test_returns_false_when_ocrmypdf_not_on_path(self, tmp_path):
|
||||
"""Returns False (and logs a warning) when ocrmypdf is not installed."""
|
||||
from app.utils.ocr_provider import embed_text_layer
|
||||
|
||||
pdf = self._make_pdf(tmp_path)
|
||||
|
||||
with patch("shutil.which", return_value=None):
|
||||
result = embed_text_layer(pdf, str(tmp_path / "out.pdf"))
|
||||
|
||||
assert result is False
|
||||
|
||||
def test_returns_true_on_success(self, tmp_path):
|
||||
"""Returns True when ocrmypdf exits with code 0."""
|
||||
from app.utils.ocr_provider import embed_text_layer
|
||||
|
||||
pdf = self._make_pdf(tmp_path)
|
||||
output = str(tmp_path / "out.pdf")
|
||||
|
||||
mock_proc = Mock()
|
||||
mock_proc.returncode = 0
|
||||
mock_proc.stderr = ""
|
||||
|
||||
with (
|
||||
patch("shutil.which", return_value="/usr/bin/ocrmypdf"),
|
||||
patch("subprocess.run", return_value=mock_proc) as mock_run,
|
||||
):
|
||||
result = embed_text_layer(pdf, output, language="eng")
|
||||
|
||||
assert result is True
|
||||
args = mock_run.call_args[0][0]
|
||||
assert "--skip-text" in args
|
||||
assert "-l" in args
|
||||
assert "eng" in args
|
||||
|
||||
def test_returns_false_on_nonzero_exit(self, tmp_path):
|
||||
"""Returns False when ocrmypdf exits with a non-zero code."""
|
||||
from app.utils.ocr_provider import embed_text_layer
|
||||
|
||||
pdf = self._make_pdf(tmp_path)
|
||||
|
||||
mock_proc = Mock()
|
||||
mock_proc.returncode = 1
|
||||
mock_proc.stderr = "some error"
|
||||
|
||||
with (
|
||||
patch("shutil.which", return_value="/usr/bin/ocrmypdf"),
|
||||
patch("subprocess.run", return_value=mock_proc),
|
||||
):
|
||||
result = embed_text_layer(pdf, str(tmp_path / "out.pdf"))
|
||||
|
||||
assert result is False
|
||||
|
||||
def test_returns_false_on_timeout(self, tmp_path):
|
||||
"""Returns False when ocrmypdf times out."""
|
||||
import subprocess
|
||||
|
||||
from app.utils.ocr_provider import embed_text_layer
|
||||
|
||||
pdf = self._make_pdf(tmp_path)
|
||||
|
||||
with (
|
||||
patch("shutil.which", return_value="/usr/bin/ocrmypdf"),
|
||||
patch("subprocess.run", side_effect=subprocess.TimeoutExpired(cmd="ocrmypdf", timeout=600)),
|
||||
):
|
||||
result = embed_text_layer(pdf, str(tmp_path / "out.pdf"))
|
||||
|
||||
assert result is False
|
||||
|
||||
def test_in_place_overwrite_on_success(self, tmp_path):
|
||||
"""When input == output the original file is replaced in-place."""
|
||||
from app.utils.ocr_provider import embed_text_layer
|
||||
|
||||
pdf = self._make_pdf(tmp_path)
|
||||
original_content = b"ORIGINAL"
|
||||
new_content = b"OCRMYPDF_OUTPUT"
|
||||
|
||||
mock_proc = Mock()
|
||||
mock_proc.returncode = 0
|
||||
mock_proc.stderr = ""
|
||||
|
||||
def fake_run(cmd: list[str], **kwargs: object) -> Mock:
|
||||
# Simulate ocrmypdf writing to the temp output path.
|
||||
out_path = cmd[-1]
|
||||
with open(out_path, "wb") as fh:
|
||||
fh.write(new_content)
|
||||
return mock_proc
|
||||
|
||||
with (
|
||||
patch("shutil.which", return_value="/usr/bin/ocrmypdf"),
|
||||
patch("subprocess.run", side_effect=fake_run),
|
||||
):
|
||||
result = embed_text_layer(pdf, pdf)
|
||||
|
||||
assert result is True
|
||||
with open(pdf, "rb") as fh:
|
||||
assert fh.read() == new_content
|
||||
|
||||
def test_in_place_cleans_up_temp_file_on_failure(self, tmp_path):
|
||||
"""Temp file created during in-place processing is removed on failure."""
|
||||
from app.utils.ocr_provider import embed_text_layer
|
||||
|
||||
pdf = self._make_pdf(tmp_path)
|
||||
|
||||
mock_proc = Mock()
|
||||
mock_proc.returncode = 1
|
||||
mock_proc.stderr = "error"
|
||||
|
||||
created_tmp: list[str] = []
|
||||
|
||||
real_mkstemp = __import__("tempfile").mkstemp
|
||||
|
||||
def fake_mkstemp(**kwargs: object) -> tuple[int, str]:
|
||||
fd, path = real_mkstemp(**kwargs)
|
||||
created_tmp.append(path)
|
||||
# Write something so the cleanup code can find the file.
|
||||
with open(path, "wb") as fh:
|
||||
fh.write(b"temp")
|
||||
return fd, path
|
||||
|
||||
with (
|
||||
patch("shutil.which", return_value="/usr/bin/ocrmypdf"),
|
||||
patch("subprocess.run", return_value=mock_proc),
|
||||
patch("tempfile.mkstemp", side_effect=fake_mkstemp),
|
||||
):
|
||||
result = embed_text_layer(pdf, pdf)
|
||||
|
||||
assert result is False
|
||||
# The temporary file should have been cleaned up.
|
||||
for tmp in created_tmp:
|
||||
assert not __import__("os").path.exists(tmp)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestProcessWithOCRTextLayerEmbedding:
|
||||
"""Tests for the embed_text_layer step in process_with_ocr."""
|
||||
|
||||
_MINIMAL_PDF = (
|
||||
b"%PDF-1.4\n"
|
||||
b"1 0 obj\n<</Type /Catalog /Pages 2 0 R>>\nendobj\n"
|
||||
b"2 0 obj\n<</Type /Pages /Kids [3 0 R] /Count 1>>\nendobj\n"
|
||||
b"3 0 obj\n<</Type /Page /Parent 2 0 R /MediaBox [0 0 612 792]>>\nendobj\n"
|
||||
b"xref\n0 4\n"
|
||||
b"0000000000 65535 f \n"
|
||||
b"0000000009 00000 n \n"
|
||||
b"0000000058 00000 n \n"
|
||||
b"0000000115 00000 n \n"
|
||||
b"trailer\n<</Size 4 /Root 1 0 R>>\n"
|
||||
b"startxref\n190\n%%EOF\n"
|
||||
)
|
||||
|
||||
@patch("app.tasks.process_with_ocr.log_task_progress")
|
||||
@patch("app.tasks.process_with_ocr.rotate_pdf_pages")
|
||||
def test_embed_text_layer_called_when_no_searchable_pdf(self, mock_rotate, mock_log, tmp_path):
|
||||
"""embed_text_layer is called when no provider returns a searchable PDF."""
|
||||
from app.tasks.process_with_ocr import process_with_ocr
|
||||
from app.utils.ocr_provider import OCRResult
|
||||
|
||||
tmp_dir = tmp_path / "tmp"
|
||||
tmp_dir.mkdir()
|
||||
pdf_file = tmp_dir / "scan.pdf"
|
||||
pdf_file.write_bytes(self._MINIMAL_PDF)
|
||||
|
||||
mock_result = OCRResult(provider="mistral", text="Hello from Mistral")
|
||||
|
||||
mock_rotate.delay = Mock()
|
||||
|
||||
with (
|
||||
patch("app.tasks.process_with_ocr.settings") as mock_settings,
|
||||
patch("app.tasks.process_with_ocr.get_ocr_providers") as mock_providers,
|
||||
patch("app.tasks.process_with_ocr.merge_ocr_results", return_value=("Hello from Mistral", None, {})),
|
||||
patch("app.tasks.process_with_ocr.embed_text_layer", return_value=True) as mock_embed,
|
||||
):
|
||||
mock_settings.workdir = str(tmp_path)
|
||||
mock_settings.tesseract_language = "eng"
|
||||
provider_mock = Mock()
|
||||
provider_mock.name = "mistral"
|
||||
provider_mock.process.return_value = mock_result
|
||||
mock_providers.return_value = [provider_mock]
|
||||
|
||||
process_with_ocr.run("scan.pdf", file_id=None)
|
||||
|
||||
mock_embed.assert_called_once_with(str(pdf_file), str(pdf_file), language="eng")
|
||||
|
||||
@patch("app.tasks.process_with_ocr.log_task_progress")
|
||||
@patch("app.tasks.process_with_ocr.rotate_pdf_pages")
|
||||
def test_embed_text_layer_skipped_when_searchable_pdf_exists(self, mock_rotate, mock_log, tmp_path):
|
||||
"""embed_text_layer is NOT called when a provider already returned a searchable PDF."""
|
||||
from app.tasks.process_with_ocr import process_with_ocr
|
||||
from app.utils.ocr_provider import OCRResult
|
||||
|
||||
tmp_dir = tmp_path / "tmp"
|
||||
tmp_dir.mkdir()
|
||||
pdf_file = tmp_dir / "scan.pdf"
|
||||
pdf_file.write_bytes(self._MINIMAL_PDF)
|
||||
|
||||
# Azure returns searchable_pdf_path set
|
||||
mock_result = OCRResult(provider="azure", text="Hello Azure", searchable_pdf_path=str(pdf_file))
|
||||
|
||||
mock_rotate.delay = Mock()
|
||||
|
||||
with (
|
||||
patch("app.tasks.process_with_ocr.settings") as mock_settings,
|
||||
patch("app.tasks.process_with_ocr.get_ocr_providers") as mock_providers,
|
||||
patch(
|
||||
"app.tasks.process_with_ocr.merge_ocr_results",
|
||||
return_value=("Hello Azure", str(pdf_file), {}),
|
||||
),
|
||||
patch("app.tasks.process_with_ocr.embed_text_layer") as mock_embed,
|
||||
):
|
||||
mock_settings.workdir = str(tmp_path)
|
||||
mock_settings.tesseract_language = "eng"
|
||||
provider_mock = Mock()
|
||||
provider_mock.name = "azure"
|
||||
provider_mock.process.return_value = mock_result
|
||||
mock_providers.return_value = [provider_mock]
|
||||
|
||||
process_with_ocr.run("scan.pdf", file_id=None)
|
||||
|
||||
mock_embed.assert_not_called()
|
||||
|
||||
@patch("app.tasks.process_with_ocr.log_task_progress")
|
||||
@patch("app.tasks.process_with_ocr.rotate_pdf_pages")
|
||||
def test_embed_text_layer_unavailable_is_handled_gracefully(self, mock_rotate, mock_log, tmp_path):
|
||||
"""When embed_text_layer returns False the task still succeeds."""
|
||||
from app.tasks.process_with_ocr import process_with_ocr
|
||||
from app.utils.ocr_provider import OCRResult
|
||||
|
||||
tmp_dir = tmp_path / "tmp"
|
||||
tmp_dir.mkdir()
|
||||
pdf_file = tmp_dir / "scan.pdf"
|
||||
pdf_file.write_bytes(self._MINIMAL_PDF)
|
||||
|
||||
mock_result = OCRResult(provider="tesseract", text="Hello Tesseract")
|
||||
mock_rotate.delay = Mock()
|
||||
|
||||
with (
|
||||
patch("app.tasks.process_with_ocr.settings") as mock_settings,
|
||||
patch("app.tasks.process_with_ocr.get_ocr_providers") as mock_providers,
|
||||
patch("app.tasks.process_with_ocr.merge_ocr_results", return_value=("Hello Tesseract", None, {})),
|
||||
patch("app.tasks.process_with_ocr.embed_text_layer", return_value=False),
|
||||
):
|
||||
mock_settings.workdir = str(tmp_path)
|
||||
mock_settings.tesseract_language = "eng+deu"
|
||||
provider_mock = Mock()
|
||||
provider_mock.name = "tesseract"
|
||||
provider_mock.process.return_value = mock_result
|
||||
mock_providers.return_value = [provider_mock]
|
||||
|
||||
result = process_with_ocr.run("scan.pdf", file_id=None)
|
||||
|
||||
# Task should succeed and return the original file path as searchable_pdf
|
||||
assert result["cleaned_text"] == "Hello Tesseract"
|
||||
assert result["searchable_pdf"] == str(pdf_file)
|
||||
|
||||
Reference in New Issue
Block a user