dsh-ab-ocr 0.1.0

package/python/ocr_worker.py

@@ -0,0 +1,516 @@
+#!/usr/bin/env python3
+"""OCR worker: one page at a time in, one Markdown document out.
+
+The worker reads a JSON job spec from stdin and writes newline-delimited JSON
+events to stdout, so the caller can report progress while a long document is
+still being recognized. Diagnostics the caller cannot act on go to stderr.
+
+The worker touches no file the caller did not name as input. It renders a page,
+recognizes it, emits that page's text and geometry so the caller can persist
+them immediately, and drops the pixels before rendering the next page. The page
+records stay in memory until the document ends, when the merge pass turns them
+into the one Markdown string and emits it on the closing event. The OCR engine
+is created for the document and released in a finally block, so a batch of
+documents never holds more than one engine.
+
+A job may instead carry a page record the caller kept from an earlier call. That
+assemble job rebuilds the pages from the record and runs the merge alone, so a
+second look at the outline costs no rendering and no recognition.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import os
+import sys
+import time
+import traceback
+from dataclasses import asdict
+from typing import Any, Sequence
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from assemble import MergeOptions, MergeStats, OutlineCandidate, build_markdown, page_text  # noqa: E402
+from clean import Page  # noqa: E402
+from layout import Line, assemble_lines  # noqa: E402
+from source import UnsupportedDocument, open_document, source_digest  # noqa: E402
+
+#: Import name and distribution name of every package the worker needs. Pillow's
+#: are different, which is why the pair is stated rather than derived.
+PACKAGES = (
+    ('rapidocr', 'rapidocr'),
+    ('onnxruntime', 'onnxruntime'),
+    ('pypdfium2', 'pypdfium2'),
+    ('PIL', 'pillow'),
+    ('numpy', 'numpy'),
+)
+
+
+def emit(event: dict[str, Any]) -> None:
+    """Write one event line to stdout as UTF-8 and flush it, so the caller sees it live.
+
+    The bytes are written to the binary stream rather than through the text
+    wrapper: a Windows console defaults to a legacy code page that cannot carry
+    the file names this tool works with.
+    """
+    sys.stdout.buffer.write((json.dumps(event, ensure_ascii=False) + '\n').encode('utf-8'))
+    sys.stdout.buffer.flush()
+
+
+def configure_stderr() -> None:
+    """Make diagnostics UTF-8, so a path in an error message is never mangled."""
+    try:
+        sys.stderr.reconfigure(encoding='utf-8', errors='backslashreplace')
+    except (AttributeError, ValueError):
+        # A stream that is not a text wrapper keeps its own settings; nothing
+        # else can reach this, because stderr is a text wrapper under CPython.
+        pass
+
+
+def read_spec(path: str | None) -> dict[str, Any]:
+    """Read the job spec the caller supplied on stdin or named with --job.
+
+    @param path: a spec file, or None to read stdin.
+    @returns: the parsed spec.
+    """
+    if path is None:
+        raw = sys.stdin.buffer.read().decode('utf-8')
+    else:
+        with open(path, 'r', encoding='utf-8') as handle:
+            raw = handle.read()
+    if raw.strip() == '':
+        raise ValueError('ocr worker: the spec carried no jobs')
+    return json.loads(raw)
+
+
+def parse_pages(spec: str | None, total: int) -> list[int]:
+    """Resolve a page selection against a document's length.
+
+    A selection is a comma-separated list of single pages and inclusive ranges,
+    counted from one, as a person would write them in a print dialog. A range
+    left open after its dash runs to the end of the document.
+    @param spec: the selection text, or None for the whole document.
+    @param total: the document's page count.
+    @returns: every selected page number, ascending and without repeats.
+    """
+    if spec is None or spec.strip() == '':
+        return list(range(1, total + 1))
+    wanted: set[int] = set()
+    for part in spec.split(','):
+        piece = part.strip()
+        if piece == '':
+            continue
+        start, end = _page_range(piece, total)
+        if start < 1 or end < start:
+            raise ValueError('invalid page selection: ' + piece)
+        wanted.update(range(start, min(end, total) + 1))
+    kept = sorted(page for page in wanted if page <= total)
+    if not kept:
+        raise ValueError('page selection names no page of a ' + str(total) + '-page document')
+    return kept
+
+
+def _page_range(piece: str, total: int) -> tuple[int, int]:
+    """Read one comma-free piece of a selection as an inclusive page range."""
+    start_text, separator, end_text = piece.partition('-')
+    start_text = start_text.strip()
+    end_text = end_text.strip()
+    if separator == '':
+        number = _page_number(piece, piece)
+        return number, number
+    start = _page_number(start_text, piece)
+    end = total if end_text == '' else _page_number(end_text, piece)
+    return start, end
+
+
+def _page_number(text: str, piece: str) -> int:
+    """Read one page number, reporting a malformed piece as a selection error."""
+    try:
+        return int(text)
+    except ValueError:
+        raise ValueError('invalid page selection: ' + piece) from None
+
+
+class Engine:
+    """The OCR engine and the recognizer call the worker makes through it."""
+
+    def __init__(self) -> None:
+        import logging
+
+        from rapidocr import RapidOCR
+
+        # The engine narrates model selection on stderr; the caller only wants
+        # failures there, and its own events travel on stdout.
+        logging.getLogger('RapidOCR').setLevel(logging.ERROR)
+        self._engine: Any = RapidOCR()
+        logging.getLogger('RapidOCR').setLevel(logging.ERROR)
+
+    def recognize(self, image: Any) -> tuple[list[Any], list[str], list[float]]:
+        """Recognize one page and return its boxes, texts, and confidences."""
+        result = self._engine(image)
+        boxes = getattr(result, 'boxes', None)
+        if boxes is None:
+            return [], [], []
+        texts = list(getattr(result, 'txts', None) or [])
+        scores = [float(score) for score in (getattr(result, 'scores', None) or [])]
+        return list(boxes), texts, scores
+
+    def release(self) -> None:
+        """Drop the engine so its ONNX sessions and model buffers are freed."""
+        self._engine = None
+        gc.collect()
+
+
+def merge_options(job: dict[str, Any]) -> MergeOptions:
+    """Build the merge pass's switches from one job's spec."""
+    merge = job.get('merge') or {}
+    return MergeOptions(
+        remove_page_numbers=bool(merge.get('removePageNumbers', True)),
+        remove_running_heads=bool(merge.get('removeRunningHeads', True)),
+        running_head_ratio=float(merge.get('runningHeadRatio', 0.6)),
+        running_head_min_pages=int(merge.get('runningHeadMinPages', 3)),
+        detect_headings=bool(merge.get('detectHeadings', True)),
+        heading_min_ratio=float(merge.get('headingMinRatio', 1.18)),
+        indent_ratio=float(merge.get('indentRatio', 1.0)),
+        paragraph_gap_ratio=float(merge.get('paragraphGapRatio', 0.85)),
+        outline_candidate_ratio=float(merge.get('outlineCandidateRatio', 1.05)),
+        max_outline_candidates=int(merge.get('maxOutlineCandidates', 400)),
+        level_overrides=level_overrides(merge.get('levelOverrides')),
+    )
+
+
+def level_overrides(value: Any) -> dict[str, Any]:
+    """Read the levels a model-assisted outline pass assigned.
+
+    @param value: the job's levelOverrides value, absent for a merge the model
+        pass did not touch.
+    @returns: the candidate id to level mapping, empty when none was sent.
+    """
+    if value is None:
+        return {}
+    if not isinstance(value, dict):
+        raise ValueError('levelOverrides is not an object')
+    return dict(value)
+
+
+def page_geometry(page: Page) -> dict[str, Any]:
+    """Describe one page so an assemble job can rebuild it without recognizing it.
+
+    @param page: the page just recognized.
+    @returns: the page's number, size, and the geometry of each of its lines.
+    """
+    return {
+        'index': page.index,
+        'width': page.width,
+        'height': page.height,
+        'lines': [
+            {
+                'text': line.text,
+                'x0': line.x0,
+                'y0': line.y0,
+                'x1': line.x1,
+                'y1': line.y1,
+                'height': line.height,
+            }
+            for line in page.lines
+        ],
+    }
+
+
+def record_pages(record: Any) -> list[Page]:
+    """Rebuild the pages an assemble job carries, in the shape the worker emits.
+
+    @param record: the job's record object.
+    @returns: the pages it lists, in its order.
+    """
+    if not isinstance(record, dict):
+        raise ValueError('assemble job carries no record object')
+    entries = record.get('pages')
+    if not isinstance(entries, list):
+        raise ValueError('assemble record carries no pages list')
+    pages: list[Page] = []
+    for entry in entries:
+        if not isinstance(entry, dict):
+            raise ValueError('assemble record carries a page that is not an object')
+        lines = entry.get('lines') or []
+        if not isinstance(lines, list):
+            raise ValueError('assemble record carries a page whose lines are not a list')
+        pages.append(Page(
+            index=int(entry['index']),
+            width=float(entry.get('width', 0.0)),
+            height=float(entry.get('height', 0.0)),
+            lines=[record_line(item) for item in lines],
+        ))
+    return pages
+
+
+def record_line(item: Any) -> Line:
+    """Rebuild one line of an assemble record from the geometry it carries.
+
+    @param item: one line object of a record page.
+    @returns: the line, with its glyph height restored from the record.
+    """
+    if not isinstance(item, dict):
+        raise ValueError('assemble record carries a line that is not an object')
+    text = item.get('text', '')
+    if not isinstance(text, str):
+        raise ValueError('assemble record carries a line whose text is not a string')
+    y0 = float(item['y0'])
+    height = float(item.get('height', float(item['y1']) - y0))
+    # A record carries geometry only: the merge reads no confidence from it.
+    return Line(
+        text=text,
+        score=1.0,
+        x0=float(item['x0']),
+        y0=y0,
+        x1=float(item['x1']),
+        y1=y0 + height,
+    )
+
+
+def merge_result(
+    job: dict[str, Any],
+    total: int,
+    recognized: int,
+    started: float,
+    markdown: str,
+    stats: MergeStats,
+    outline: Sequence[OutlineCandidate],
+    truncated: bool,
+) -> dict[str, Any]:
+    """Build the closing event's fields for one merged document.
+
+    @param job: the job the result belongs to.
+    @param total: the document's page count.
+    @param recognized: how many pages the job recognized or carried.
+    @param started: when the job began, so the result reports its own cost.
+    @param markdown: the merged document.
+    @param stats: the merge pass's counts.
+    @param outline: the candidates a model-assisted pass may re-level.
+    @param truncated: whether the candidate list hit its cap.
+    @returns: the fields of the job's done event.
+    """
+    return {
+        'id': job.get('id', ''),
+        'input': job.get('input', ''),
+        'pages': stats.pages,
+        'totalPages': total,
+        'lines': stats.lines,
+        'headings': stats.headings,
+        'chars': len(markdown),
+        'droppedPageNumbers': stats.dropped_page_numbers,
+        'droppedRunningHeads': stats.dropped_running_heads,
+        'joinedAcrossPages': stats.joined_across_pages,
+        'recognized': recognized,
+        'seconds': round(time.time() - started, 2),
+        'markdown': markdown,
+        'outline': [asdict(candidate) for candidate in outline],
+        'outlineTruncated': truncated,
+    }
+
+
+def failure(job: dict[str, Any], error: BaseException) -> dict[str, Any]:
+    """One failed document, in the terms the caller reports it with."""
+    return {'id': job.get('id', ''), 'input': job.get('input', ''), 'error': type(error).__name__ + ': ' + str(error)}
+
+
+def run_job(
+    job: dict[str, Any],
+    engine_lifetime: str,
+    shared: Engine | None,
+) -> tuple[dict[str, Any] | None, dict[str, Any] | None, Engine | None]:
+    """Recognize and merge one document, or merge one the caller already read.
+
+    @param job: one entry of the spec's job list.
+    @param engine_lifetime: 'perDocument' or 'shared'.
+    @param shared: the engine a shared lifetime carries between documents.
+    @returns: the document result or the failure, and the engine still in use.
+    """
+    if str(job.get('mode', 'recognize')) == 'assemble':
+        try:
+            return assemble_job(job), None, shared
+        except Exception as error:  # noqa: BLE001 - one unusable record must not end the batch
+            traceback.print_exc(file=sys.stderr)
+            return None, failure(job, error), shared
+
+    path = job['input']
+    started = time.time()
+    try:
+        document = open_document(path, int(job.get('dpi', 200)), int(job.get('maxPixels', 0)))
+    except UnsupportedDocument as error:
+        return None, failure(job, error), shared
+    except Exception as error:  # noqa: BLE001 - the caller reports any open failure
+        return None, failure(job, error), shared
+
+    total = 0
+    engine = shared
+    recognized = 0
+    pages: list[Page] = []
+    try:
+        total = document.page_count()
+        max_pages = int(job.get('maxPages', 0))
+        if max_pages > 0 and total > max_pages and not str(job.get('pages') or '').strip():
+            raise ValueError(
+                'document carries ' + str(total) + ' pages, above the ' + str(max_pages)
+                + '-page ceiling; select pages or raise maxPages'
+            )
+        selection = parse_pages(job.get('pages'), total)
+        # The caller names a recognition after this digest, so it has to describe
+        # the document's content rather than its length: a source edited without
+        # changing size is a different recognition and must not reuse the name.
+        emit({
+            'event': 'start',
+            'id': job.get('id', ''),
+            'input': path,
+            'pages': selection,
+            'totalPages': total,
+            'sourceDigest': source_digest(path),
+        })
+        if engine is None:
+            engine = Engine()
+        for number in selection:
+            rendered = document.render(number - 1)
+            boxes, texts, scores = engine.recognize(rendered.image)
+            lines = assemble_lines(
+                boxes,
+                texts,
+                scores,
+                float(rendered.width),
+                float(job.get('textScore', 0.5)),
+                bool(job.get('detectColumns', True)),
+            )
+            page = Page(index=number, width=float(rendered.width), height=float(rendered.height), lines=lines)
+            pages.append(page)
+            recognized += 1
+            emit({
+                'event': 'page',
+                'id': job.get('id', ''),
+                'page': page.index,
+                'totalPages': total,
+                'lines': len(lines),
+                'chars': sum(len(line.text) for line in lines),
+                'text': page_text(lines),
+                'geometry': page_geometry(page),
+            })
+            # Dropping the page before the next render is what bounds memory.
+            del rendered, boxes, texts, scores, lines, page
+    except Exception as error:  # noqa: BLE001 - one bad document must not end the batch
+        traceback.print_exc(file=sys.stderr)
+        if engine_lifetime == 'perDocument' and engine is not None:
+            engine.release()
+        return None, failure(job, error), None if engine_lifetime == 'perDocument' else shared
+    finally:
+        document.close()
+
+    if engine_lifetime == 'perDocument' and engine is not None:
+        engine.release()
+        engine = None
+        emit({'event': 'released', 'id': job.get('id', ''), 'reason': 'document-finished'})
+
+    try:
+        markdown, stats, outline, truncated = build_markdown(pages, merge_options(job))
+    except Exception as error:  # noqa: BLE001 - a bad merge spec is one failed document
+        traceback.print_exc(file=sys.stderr)
+        return None, failure(job, error), None if engine_lifetime == 'perDocument' else shared
+    pages.clear()
+    result = merge_result(job, total, recognized, started, markdown, stats, outline, truncated)
+    return result, None, shared if engine_lifetime == 'shared' else engine
+
+
+def assemble_job(job: dict[str, Any]) -> dict[str, Any]:
+    """Merge one document whose pages the caller already recognized.
+
+    The caller keeps the page records a recognize job emitted, so a second look
+    at the outline costs the merge and nothing else: no document is opened and
+    no engine is built.
+    @param job: one entry of the spec's job list, carrying a page record.
+    @returns: the fields of the job's done event.
+    """
+    started = time.time()
+    pages = record_pages(job.get('record'))
+    markdown, stats, outline, truncated = build_markdown(pages, merge_options(job))
+    total = len(pages)
+    pages.clear()
+    return merge_result(job, total, total, started, markdown, stats, outline, truncated)
+
+
+def self_test() -> int:
+    """Report the worker's environment as one event, for the caller's check.
+
+    Importing the packages is the check: a missing or mismatched dependency is
+    what a fresh installation gets wrong, and it costs a second to detect. The
+    engine itself is built per document, where a failure names the document.
+    """
+    report: dict[str, Any] = {'event': 'ready', 'python': sys.version.split()[0], 'modules': {}}
+    missing: list[str] = []
+    for module, distribution in PACKAGES:
+        try:
+            __import__(module)
+        except Exception as error:  # noqa: BLE001 - a missing package is the answer
+            report['modules'][module] = None
+            report['modules'][module + 'Error'] = type(error).__name__ + ': ' + str(error)
+            missing.append(module)
+            continue
+        try:
+            from importlib.metadata import version
+
+            report['modules'][module] = version(distribution)
+        except Exception:  # noqa: BLE001 - an importable package with no metadata is still present
+            report['modules'][module] = 'unknown'
+    report['ready'] = not missing
+    emit(report)
+    return 0 if not missing else 3
+
+
+def main(argv: Sequence[str]) -> int:
+    """Read the spec, process every job, and report the batch outcome."""
+    parser = argparse.ArgumentParser(description='OCR a PDF or image into Markdown.')
+    parser.add_argument('--job', help='path to a JSON spec instead of stdin')
+    parser.add_argument('--self-test', action='store_true', help='report the environment and exit')
+    args = parser.parse_args(argv)
+    configure_stderr()
+
+    if args.self_test:
+        return self_test()
+
+    try:
+        spec = read_spec(args.job)
+    except Exception as error:  # noqa: BLE001 - a bad spec is reported, not raised
+        emit({'event': 'fatal', 'error': type(error).__name__ + ': ' + str(error)})
+        return 2
+
+    jobs = spec.get('jobs', [])
+    lifetime = str(spec.get('engineLifetime', 'perDocument'))
+    shared: Engine | None = None
+    # A batch of assemble jobs merges records the caller already has, so a
+    # shared lifetime costs no engine when no job recognizes.
+    if lifetime == 'shared' and any(str(job.get('mode', 'recognize')) != 'assemble' for job in jobs):
+        try:
+            shared = Engine()
+        except Exception as error:  # noqa: BLE001 - report an unusable engine once
+            emit({'event': 'fatal', 'error': type(error).__name__ + ': ' + str(error)})
+            return 3
+
+    finished = 0
+    failures: list[dict[str, Any]] = []
+    for job in jobs:
+        result, problem, shared = run_job(job, lifetime, shared)
+        if result is not None:
+            finished += 1
+            emit({'event': 'done', **result})
+        if problem is not None:
+            failures.append(problem)
+            emit({'event': 'error', 'id': problem['id'], 'input': problem['input'], 'message': problem['error']})
+
+    if shared is not None:
+        shared.release()
+        emit({'event': 'released', 'id': '', 'reason': 'batch-finished'})
+
+    emit({'event': 'end', 'documents': finished, 'failures': failures})
+    return 0 if finished > 0 else 1
+
+
+if __name__ == '__main__':
+    sys.exit(main(sys.argv[1:]))

package/python/requirements.txt

@@ -0,0 +1,16 @@
+# Runtime dependencies of the OCR worker.
+#
+# rapidocr ships the PP-OCR detection, classification, and recognition models
+# inside the wheel, so recognition works offline once these are installed and
+# "installing the model" is the same act as installing the package.
+# onnxruntime is the inference engine rapidocr selects by default; rapidocr
+# imports numpy directly, so it is stated here rather than relied on as a
+# transitive dependency of either.
+# pypdfium2 renders PDF pages and carries its own PDFium build, so no system
+# PDF library is needed.
+# pillow decodes still images, including every frame of a multi-page one.
+rapidocr>=3.0,<4
+onnxruntime>=1.17
+numpy>=1.26
+pypdfium2>=4.0
+pillow>=10.0