feat(comments): extract-ocr — RapidOCR pass for scanned attachments (refs #663)
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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}
|
||||
)
|
||||
|
||||
@@ -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
51
src/rex/comments/ocr.py
Normal 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)
|
||||
@@ -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):
|
||||
|
||||
82
tests/rex/comments/test_extract_ocr.py
Normal file
82
tests/rex/comments/test_extract_ocr.py
Normal 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
|
||||
Reference in New Issue
Block a user