feat(comments): extract-ocr — RapidOCR pass for scanned attachments (refs #663)

This commit is contained in:
kert
2026-08-20 10:32:44 -04:00
parent 1d1bdee84f
commit d94fbb36ea
6 changed files with 238 additions and 19 deletions

View File

@@ -157,16 +157,74 @@ def extract(
@app.command(name="extract-ocr")
def extract_ocr(
root: Path = typer.Option(_DEFAULT_ROOT, "--root"), # noqa: ARG001
root: Path = typer.Option(_DEFAULT_ROOT, "--root", help="Comments root dir."),
docket: str = typer.Option(
None, "--docket", help="Limit to one docket (e.g. CMS-2023-0121)."
),
limit: int = typer.Option(
0, "--limit", help="Cap comment dirs processed this run. 0 = no cap."
),
) -> None:
"""OCR pass for ocr_needed attachments. Phase 2 — not yet implemented."""
typer.echo(
"extract-ocr is phase 2 — not yet implemented. "
"Tracks: see #253 follow-up. "
"Today, ocr_needed attachments stay flagged in combined.md frontmatter.",
err=True,
)
raise typer.Exit(2)
"""OCR pass for ``ocr_needed`` attachments (#663).
Finds comment dirs whose combined.md frontmatter flags any
attachment ``ocr_needed`` and re-extracts them with the RapidOCR
engine (rewrites combined.md + sibling transcripts, keeps the flag
on anything OCR can't clear). CPU-only, pip-only:
uv run --with rapidocr-onnxruntime stack comments extract-ocr
"""
import yaml
from rex.comments.combine import extract_comment
from rex.comments.ocr import RapidOcrEngine
if not root.is_dir():
typer.echo(f"root not found: {root}", err=True)
raise typer.Exit(1)
pending: list[Path] = []
dockets = sorted(root.glob(f"{docket or 'CMS-'}*"))
for docket_dir in dockets:
if not docket_dir.is_dir():
continue
for combined in docket_dir.glob("*/combined.md"):
text = combined.read_text(encoding="utf-8", errors="replace")
if "ocr_needed" not in text:
continue
try:
end = text.find("\n---\n", 4)
fm = yaml.safe_load(text[4:end]) or {}
except Exception: # noqa: BLE001 — malformed frontmatter: retry it
fm = {}
atts = fm.get("attachments") or []
if atts and not any(a.get("status") == "ocr_needed" for a in atts):
continue
pending.append(combined.parent)
if limit:
pending = pending[: int(limit)]
typer.echo(f"ocr queue: {len(pending)} comment dirs")
if not pending:
return
body_lookup = _bib_lookup_factory(True)
engine = RapidOcrEngine()
done = errors = 0
for i, comment_dir in enumerate(pending, 1):
try:
extract_comment(
comment_dir,
inline_body=body_lookup(comment_dir.name),
force=True,
ocr_engine=engine,
)
done += 1
except Exception as e: # noqa: BLE001
log.warning("ocr extract failed for %s: %s", comment_dir.name, e)
errors += 1
if i % 25 == 0:
typer.echo(f" {i}/{len(pending)} ok={done} errors={errors}", err=False)
typer.echo(f" done: {done} re-extracted, {errors} errors")
@app.command()

View File

@@ -24,6 +24,7 @@ def extract_comment(
*,
inline_body: str = "",
force: bool = False,
ocr_engine=None,
) -> Path:
"""Extract every attachment in *comment_dir* and write *combined.md*.
@@ -49,7 +50,7 @@ def extract_comment(
sibling_payloads: list[tuple[str, str]] = [] # (sibling_filename, contents)
for path in sorted(_attachment_paths(comment_dir)):
result = extract_attachment(path)
result = extract_attachment(path, ocr_engine=ocr_engine)
attachments_meta.append(
{"file": path.name, "status": result.status, "chars": result.chars}
)

View File

@@ -29,11 +29,18 @@ class ExtractResult:
_IMAGE_SUFFIXES = (".jpg", ".jpeg", ".png", ".tif", ".tiff", ".gif", ".bmp")
def extract_attachment(path: Path) -> ExtractResult:
"""Extract text from one attachment. See module docstring for status values."""
def extract_attachment(path: Path, *, ocr_engine=None) -> ExtractResult:
"""Extract text from one attachment. See module docstring for status values.
``ocr_engine`` is an optional callable ``(path) -> str`` used when a
file would otherwise land in the ``ocr_needed`` queue (scanned
images, PDFs without a text layer). Its output replaces the empty
extraction when it clears the OCR threshold (#663).
"""
suffix = path.suffix.lower()
if suffix == ".pdf":
return _extract_pdf(path)
result = _extract_pdf(path)
return _maybe_ocr(path, result, ocr_engine)
if suffix == ".docx":
return _extract_docx(path)
if suffix in (".txt", ".html", ".htm"):
@@ -41,10 +48,30 @@ def extract_attachment(path: Path) -> ExtractResult:
if suffix in _IMAGE_SUFFIXES:
# Real CMS comments include scanned-letter images; route to
# phase-2 OCR queue, don't drop them as "unsupported".
return ExtractResult(text="", status="ocr_needed", chars=0)
return _maybe_ocr(
path, ExtractResult(text="", status="ocr_needed", chars=0), ocr_engine
)
return ExtractResult(text="", status="unsupported", chars=0)
def _maybe_ocr(path: Path, result: ExtractResult, ocr_engine) -> ExtractResult:
"""Run *ocr_engine* over files the plain extractors flagged ocr_needed."""
if result.status != "ocr_needed" or ocr_engine is None:
return result
try:
text = ocr_engine(path) or ""
except Exception as e: # noqa: BLE001 — engine failure keeps the queue flag
log.warning("ocr failed for %s: %s", path, e)
return result
if result.text.strip():
# keep whatever text layer the PDF did have, OCR appended
text = result.text + "\n\n" + text
chars = len(text)
if chars <= _OCR_THRESHOLD:
return result
return ExtractResult(text=text, status="ok", chars=chars)
def _extract_pdf(path: Path) -> ExtractResult:
try:
with fitz.open(path) as doc:

51
src/rex/comments/ocr.py Normal file
View File

@@ -0,0 +1,51 @@
"""RapidOCR-backed engine for scanned comment attachments (#663).
The engine is a callable ``(path) -> str`` matching the ``ocr_engine``
hook in :mod:`rex.comments.extract`. Images OCR directly; PDFs render
per page through PyMuPDF (no poppler dependency) at 200 dpi.
RapidOCR (onnxruntime backend) is pip-only — no tesseract or system
packages. It is NOT in the base environment: run via
uv run --with rapidocr-onnxruntime stack comments extract-ocr
"""
from __future__ import annotations
import logging
from pathlib import Path
import fitz # pymupdf
log = logging.getLogger(__name__)
_PDF_MAX_PAGES = 40 # scanned letters, not books — cap runaway PDFs
_DPI = 200
class RapidOcrEngine:
"""Lazy RapidOCR wrapper; one model load per process."""
def __init__(self) -> None:
from rapidocr_onnxruntime import RapidOCR
self._ocr = RapidOCR()
def __call__(self, path: Path) -> str:
if path.suffix.lower() == ".pdf":
return self._pdf(path)
return self._image_bytes(path.read_bytes())
def _image_bytes(self, data: bytes) -> str:
result, _ = self._ocr(data)
return " ".join(r[1] for r in (result or []))
def _pdf(self, path: Path) -> str:
parts: list[str] = []
with fitz.open(path) as doc:
for page in doc[:_PDF_MAX_PAGES]:
pix = page.get_pixmap(dpi=_DPI)
text = self._image_bytes(pix.tobytes("png"))
if text:
parts.append(text)
return "\n\n".join(parts)

View File

@@ -40,12 +40,12 @@ def test_extract_writes_combined(tmp_path: Path):
assert "written" in result.output.lower()
def test_extract_ocr_is_stub(tmp_path: Path):
def test_extract_ocr_empty_queue(tmp_path: Path):
# Implemented in #663 — an empty root reports a zero queue and exits
# cleanly without loading the OCR engine or the bib lookup.
result = runner.invoke(app, ["extract-ocr", "--root", str(tmp_path)])
assert result.exit_code != 0
assert (
"phase 2" in result.output.lower() or "not implemented" in result.output.lower()
)
assert result.exit_code == 0
assert "ocr queue: 0" in result.output
def test_stats_after_extract(tmp_path: Path):

View File

@@ -0,0 +1,82 @@
"""Tests for OCR-augmented extraction (rex/comments — #663)."""
from __future__ import annotations
from pathlib import Path
import fitz
from rex.comments.combine import extract_comment, parse_combined
from rex.comments.extract import extract_attachment
def _fake_engine(text: str):
def engine(path: Path) -> str: # noqa: ARG001
return text
return engine
def _write_image(path: Path) -> None:
# 1x1 PNG via pymupdf — content is irrelevant, the fake engine answers.
pix = fitz.Pixmap(fitz.csRGB, fitz.IRect(0, 0, 8, 8))
pix.save(str(path))
def _write_scanned_pdf(path: Path) -> None:
# a PDF with no text layer
doc = fitz.open()
doc.new_page()
doc.save(str(path))
doc.close()
class TestExtractAttachmentOcr:
def test_image_without_engine_stays_queued(self, tmp_path):
img = tmp_path / "a.png"
_write_image(img)
r = extract_attachment(img)
assert r.status == "ocr_needed"
def test_image_with_engine_extracts(self, tmp_path):
img = tmp_path / "a.png"
_write_image(img)
r = extract_attachment(img, ocr_engine=_fake_engine("Dear CMS, " * 20))
assert r.status == "ok"
assert "Dear CMS" in r.text
def test_scanned_pdf_with_engine_extracts(self, tmp_path):
pdf = tmp_path / "scan.pdf"
_write_scanned_pdf(pdf)
r = extract_attachment(
pdf, ocr_engine=_fake_engine("Scanned letter text " * 10)
)
assert r.status == "ok"
assert "Scanned letter" in r.text
def test_engine_short_output_stays_queued(self, tmp_path):
img = tmp_path / "a.png"
_write_image(img)
r = extract_attachment(img, ocr_engine=_fake_engine("x"))
assert r.status == "ocr_needed"
class TestExtractCommentOcrRerun:
def test_force_rerun_with_engine_updates_combined(self, tmp_path):
d = tmp_path / "CMS-2023-0121" / "CMS-2023-0121-9999"
d.mkdir(parents=True)
_write_image(d / "attachment_1.png")
# first pass without OCR → ocr_needed in frontmatter
out = extract_comment(d, inline_body="inline")
fm, _ = parse_combined(out.read_text())
assert fm["attachments"][0]["status"] == "ocr_needed"
# OCR pass
out2 = extract_comment(
d,
inline_body="inline",
force=True,
ocr_engine=_fake_engine("Recovered scanned palliative text " * 5),
)
fm2, body2 = parse_combined(out2.read_text())
assert fm2["attachments"][0]["status"] == "ok"
assert "Recovered scanned palliative text" in body2