copilot-swe-agent[bot]
0b291995b9
- Fix _get_pipeline_ocr_language: remove redundant `or None` in step_config.get() - Add Session type hint to _get_pipeline_ocr_language db parameter via TYPE_CHECKING - Update process_with_ocr to use modern str | None syntax instead of Optional[str] - Fix test_get_pipeline_ocr_language_explicit_pipeline_takes_priority: properly add sys_step to db_session so the system pipeline step is persisted in the test DB Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
304 lines
13 KiB
Python
304 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
"""Unified OCR processing task for DocuElevate.
|
||
|
||
This task replaces the single-provider ``process_with_azure_document_intelligence``
|
||
task with a multi-engine OCR pipeline that:
|
||
|
||
1. Runs every OCR provider listed in ``OCR_PROVIDERS`` (default: ``azure``).
|
||
2. Merges/cross-checks the results using the configured AI model when more
|
||
than one provider is active (see ``OCR_MERGE_STRATEGY``).
|
||
3. Writes the best searchable PDF back to the working directory.
|
||
4. Optionally compares the OCR output against the original embedded text
|
||
(passed as *original_text*) using a head-to-head AI review and keeps the
|
||
higher-quality text for downstream processing.
|
||
5. Hands off to the page-rotation and metadata-extraction pipeline exactly as
|
||
the legacy Azure task did.
|
||
"""
|
||
|
||
import logging
|
||
import os
|
||
|
||
from app.celery_app import celery
|
||
from app.config import settings
|
||
from app.database import SessionLocal
|
||
from app.models import FileRecord
|
||
from app.tasks.retry_config import OcrTaskWithRetry
|
||
from app.tasks.rotate_pdf_pages import rotate_pdf_pages
|
||
from app.utils import log_task_progress
|
||
from app.utils.ocr_provider import OCRResult, embed_text_layer, get_ocr_providers, merge_ocr_results
|
||
from app.utils.text_quality import TextSource, check_text_quality, compare_text_quality
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@celery.task(base=OcrTaskWithRetry, bind=True)
|
||
def process_with_ocr(
|
||
self,
|
||
filename: str,
|
||
file_id: int | None = None,
|
||
original_text: str | None = None,
|
||
language: str | None = None,
|
||
):
|
||
"""Run the configured OCR providers on *filename* and continue the pipeline.
|
||
|
||
When multiple OCR providers are configured the results are merged using the
|
||
AI model (or a simpler strategy controlled by ``OCR_MERGE_STRATEGY``).
|
||
|
||
If *original_text* is provided (the original embedded text that failed the
|
||
quality check), the OCR result is compared against it using a head-to-head
|
||
AI review. The higher-quality text is passed to downstream tasks.
|
||
|
||
Args:
|
||
filename: Base name of the file inside ``<workdir>/tmp/``.
|
||
file_id: Optional database record ID passed through to downstream tasks.
|
||
original_text: Optional original embedded text for head-to-head comparison.
|
||
language: Optional Tesseract-style language code(s) (e.g. ``"eng+deu"``)
|
||
to override the global OCR language settings for this specific run.
|
||
Pass ``None`` or ``"auto"`` to use the global settings. This
|
||
enables per-pipeline language configuration.
|
||
"""
|
||
task_id = self.request.id
|
||
log_task_progress(
|
||
task_id,
|
||
"process_with_ocr",
|
||
"in_progress",
|
||
f"Starting OCR for {filename}",
|
||
file_id=file_id,
|
||
)
|
||
|
||
try:
|
||
tmp_file_path = os.path.join(settings.workdir, "tmp", filename)
|
||
if not os.path.exists(tmp_file_path):
|
||
raise FileNotFoundError(f"Local file not found: {tmp_file_path}")
|
||
|
||
providers = get_ocr_providers(language=language)
|
||
provider_names = [p.name for p in providers]
|
||
logger.info(f"[{task_id}] Running {len(providers)} OCR provider(s): {provider_names}")
|
||
|
||
log_task_progress(
|
||
task_id,
|
||
"run_ocr_providers",
|
||
"in_progress",
|
||
f"Running OCR providers: {', '.join(provider_names)}",
|
||
file_id=file_id,
|
||
)
|
||
|
||
results = []
|
||
errors = []
|
||
for provider in providers:
|
||
pname = provider.__class__.__name__
|
||
try:
|
||
logger.info(f"[{task_id}] Running {pname} on {filename}")
|
||
result: OCRResult = provider.process(tmp_file_path)
|
||
results.append(result)
|
||
logger.info(f"[{task_id}] {pname} extracted {len(result.text)} chars")
|
||
except Exception as exc:
|
||
logger.error(f"[{task_id}] {pname} failed for {filename}: {exc}")
|
||
errors.append(f"{pname}: {exc}")
|
||
|
||
if not results:
|
||
error_summary = "; ".join(errors)
|
||
log_task_progress(
|
||
task_id,
|
||
"run_ocr_providers",
|
||
"failure",
|
||
"All OCR providers failed",
|
||
file_id=file_id,
|
||
detail=error_summary,
|
||
)
|
||
raise RuntimeError(f"All OCR providers failed for {filename}: {error_summary}")
|
||
|
||
if errors:
|
||
logger.warning(f"[{task_id}] Some OCR providers failed: {'; '.join(errors)}")
|
||
|
||
log_task_progress(
|
||
task_id,
|
||
"run_ocr_providers",
|
||
"success",
|
||
f"{len(results)} of {len(providers)} OCR provider(s) succeeded",
|
||
file_id=file_id,
|
||
)
|
||
|
||
# Merge results (no-op when only one provider succeeded)
|
||
extracted_text, searchable_pdf_path, rotation_data = merge_ocr_results(results, filename)
|
||
logger.info(
|
||
f"[{task_id}] Merged OCR text: {len(extracted_text)} chars, "
|
||
f"pdf={'yes' if searchable_pdf_path else 'no'}, "
|
||
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:
|
||
# Use the per-call language override; fall back to global setting
|
||
embed_lang = (
|
||
language
|
||
if language and language != "auto"
|
||
else (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=embed_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,
|
||
)
|
||
|
||
# ----------------------------------------------------------------
|
||
# Head-to-head comparison with original embedded text (if provided)
|
||
# ----------------------------------------------------------------
|
||
final_text = extracted_text
|
||
if original_text and original_text.strip() and extracted_text.strip():
|
||
logger.info(f"[{task_id}] Original embedded text provided; running head-to-head quality comparison")
|
||
log_task_progress(
|
||
task_id,
|
||
"compare_ocr_quality",
|
||
"in_progress",
|
||
"Comparing OCR result against original embedded text",
|
||
file_id=file_id,
|
||
)
|
||
try:
|
||
comparison = compare_text_quality(original_text, extracted_text)
|
||
comparison_detail = (
|
||
f"Original score: {comparison.original_score}/100, "
|
||
f"OCR score: {comparison.ocr_score}/100, "
|
||
f"Preferred: {comparison.preferred}\n"
|
||
f"AI explanation: {comparison.explanation}"
|
||
)
|
||
logger.info(f"[{task_id}] OCR comparison – {comparison_detail}")
|
||
|
||
if comparison.preferred == "original":
|
||
# Original text is actually better – use it instead of OCR.
|
||
final_text = original_text
|
||
logger.info(
|
||
f"[{task_id}] Original embedded text selected "
|
||
f"(original={comparison.original_score} > ocr={comparison.ocr_score})"
|
||
)
|
||
log_task_progress(
|
||
task_id,
|
||
"compare_ocr_quality",
|
||
"success",
|
||
f"Original text preferred (original={comparison.original_score}/100 vs "
|
||
f"ocr={comparison.ocr_score}/100)",
|
||
file_id=file_id,
|
||
detail=comparison_detail,
|
||
)
|
||
else:
|
||
logger.info(
|
||
f"[{task_id}] OCR text selected "
|
||
f"(preferred={comparison.preferred!r}, "
|
||
f"ocr={comparison.ocr_score}, original={comparison.original_score})"
|
||
)
|
||
log_task_progress(
|
||
task_id,
|
||
"compare_ocr_quality",
|
||
"success",
|
||
f"OCR text preferred (ocr={comparison.ocr_score}/100 vs "
|
||
f"original={comparison.original_score}/100)",
|
||
file_id=file_id,
|
||
detail=comparison_detail,
|
||
)
|
||
except Exception as cmp_exc:
|
||
logger.warning(f"[{task_id}] Head-to-head comparison failed ({cmp_exc}); keeping OCR text")
|
||
log_task_progress(
|
||
task_id,
|
||
"compare_ocr_quality",
|
||
"skipped",
|
||
f"Comparison failed ({cmp_exc}); keeping OCR output",
|
||
file_id=file_id,
|
||
)
|
||
elif original_text is not None:
|
||
# original_text was provided but one side is empty – pick whichever has content.
|
||
if not extracted_text.strip() and original_text.strip():
|
||
final_text = original_text
|
||
logger.info(f"[{task_id}] OCR returned empty text; falling back to original embedded text")
|
||
log_task_progress(
|
||
task_id,
|
||
"compare_ocr_quality",
|
||
"success",
|
||
"OCR empty – using original embedded text",
|
||
file_id=file_id,
|
||
)
|
||
else:
|
||
log_task_progress(
|
||
task_id,
|
||
"compare_ocr_quality",
|
||
"skipped",
|
||
"No original text to compare; using OCR output",
|
||
file_id=file_id,
|
||
)
|
||
|
||
log_task_progress(
|
||
task_id,
|
||
"process_with_ocr",
|
||
"success",
|
||
f"OCR complete for {filename}",
|
||
file_id=file_id,
|
||
detail=f"Extracted {len(extracted_text)} chars using {len(results)} provider(s); "
|
||
f"final text length: {len(final_text)} chars",
|
||
)
|
||
|
||
# Score the final embedded text — the text that will land in ocr_text.
|
||
# We always call check_text_quality() on final_text because:
|
||
# - merge_ocr_results() may have AI-merged output from several engines
|
||
# - compare_text_quality() scores are relative (not the same scale)
|
||
# - The original may have been preferred, reversing the OCR output
|
||
# OCR-produced (or AI-merged) text is treated as TextSource.OCR_PREVIOUS
|
||
# so the quality AI call is always made.
|
||
if file_id is not None:
|
||
try:
|
||
quality_result = check_text_quality(final_text, TextSource.OCR_PREVIOUS)
|
||
logger.info(
|
||
f"[{task_id}] Final text quality: score={quality_result.quality_score}/100, "
|
||
f"good={quality_result.is_good_quality}, feedback={quality_result.feedback!r}"
|
||
)
|
||
with SessionLocal() as _db:
|
||
_rec = _db.query(FileRecord).filter_by(id=file_id).first()
|
||
if _rec:
|
||
_rec.ocr_quality_score = quality_result.quality_score
|
||
_db.commit()
|
||
logger.info(f"[{task_id}] Saved ocr_quality_score={quality_result.quality_score} for file_id={file_id}")
|
||
except Exception as _score_exc:
|
||
logger.warning(f"[{task_id}] Could not persist ocr_quality_score: {_score_exc}")
|
||
|
||
# Continue pipeline: rotate pages (if needed), then extract metadata
|
||
rotate_pdf_pages.delay(filename, final_text, rotation_data, file_id)
|
||
|
||
return {
|
||
"file": filename,
|
||
"searchable_pdf": searchable_pdf_path or tmp_file_path,
|
||
"cleaned_text": final_text,
|
||
"providers_used": [r.provider for r in results],
|
||
}
|
||
|
||
except Exception as exc:
|
||
logger.error(f"[{task_id}] OCR failed for {filename}: {exc}")
|
||
log_task_progress(
|
||
task_id,
|
||
"process_with_ocr",
|
||
"failure",
|
||
f"OCR failed for {filename}",
|
||
file_id=file_id,
|
||
detail=str(exc),
|
||
)
|
||
raise
|