Download runtime/jevlike/ingest.py from Cem13/kodama-core: direct link, hf CLI and curl.
- Browser
- Download file 55.3 kB
-
https://huggingface.co/Cem13/kodama-core/resolve/main/runtime/jevlike/ingest.py
- Command line
-
hf download hf://Cem13/kodama-core/runtime/jevlike/ingest.py
-
curl -L -o ingest.py https://huggingface.co/Cem13/kodama-core/resolve/main/runtime/jevlike/ingest.py
55.3 kB
| """Document ingestion: PDFs, scans, photos, HTML and text -> clean text for SystemOne. | |
| The model reads text only, so every other input is converted first: | |
| from jevlike.ingest import to_text, predict_file | |
| r = to_text("10-K.pdf") # IngestResult(text, pages, method_per_page, seconds, warnings, ...) | |
| r = to_text("receipt.jpg", ocr="rapidocr") | |
| r = to_text(["p1.png", "p2.png"]) # one document given as several page images | |
| r = to_text(pdf_bytes) # bytes: the kind is sniffed from the content | |
| res = predict_file(model, "10-K.pdf", {"topic": Choice(...), ...}) # ingest + SystemOne.predict | |
| python -m jevlike.ingest FILE [--ocr rapidocr] [--json] [--questions spec.json --checkpoint DIR] | |
| What each input goes through | |
| ---------------------------- | |
| * **PDF**: the text layer of every page via PyMuPDF (content-stream order, which keeps the columns | |
| of multi-column layouts apart). Ruled tables found by pdfplumber are replaced by markdown rows | |
| (``| Revenue | 1,234 | 1,456 |``) so numbers stay next to their row and column headers. A page | |
| with too little extractable text (< ``min_chars`` non-space characters, mostly unmapped glyphs, | |
| or a page-sized image with only a stamp of text) is rendered at ``dpi`` (300) and OCR'd. | |
| * **Images** (.png .jpg .jpeg .tif/.tiff (all frames) .webp .bmp .gif): OCR with tesseract | |
| (default) or RapidOCR (``ocr="rapidocr"``, PP-OCRv4 via onnxruntime). Before OCR | |
| (``preprocess=True``): EXIF orientation, grayscale, perspective correction of a photographed | |
| page (largest bright quadrilateral on a darker background), a 3x3 median filter when the paper | |
| shows pixel noise, illumination flattening when the background is uneven, deskew (+-5 deg), and | |
| upscaling when the text is small (x-height < ~20 px; tesseract only). RapidOCR's bundled Chinese | |
| PP-OCRv4 recognizer drops spaces between English words; set ``JEVLIKE_RAPIDOCR_REC`` to an | |
| English recognizer (``en_PP-OCRv3_rec_infer.onnx`` from the ``SWHL/RapidOCR`` HF repo, 9 MB). | |
| ``JEVLIKE_OCR_THREADS`` limits RapidOCR's onnxruntime threads. | |
| * **HTML**: scripts, styles, navigation, headers/footers outside the article, sidebars, cookie and | |
| share widgets are dropped; the ``<article>`` / ``<main>`` content is kept as markdown-ish text | |
| (``#`` headings, paragraphs, ``-`` list items, ``|`` table rows). | |
| * **Text / Markdown**: decoded (UTF-8, else cp1252/latin-1) and whitespace-normalized. | |
| Layout handling shared by the text-layer and OCR paths: side-by-side narrow blocks whose lines | |
| line up (a borderless table that OCR / the PDF writer emits column by column) are zipped back | |
| into rows; cells on one line separated by a wide gap are joined with `` | ``; lines repeated at | |
| the top or bottom of most pages (running headers, "Page 3 of 9") are removed from multi-page | |
| documents. Cleanup: ligatures / soft hyphens / zero-width characters, de-hyphenation of words | |
| split across lines, re-flowing of wrapped lines into paragraphs (table rows and list items are | |
| kept as lines), OCR noise lines without any letter or digit, whitespace normalization. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import inspect | |
| import io | |
| import json | |
| import os | |
| import re | |
| import shutil | |
| import statistics | |
| import subprocess | |
| import sys | |
| import time | |
| import unicodedata | |
| from dataclasses import asdict, dataclass, field | |
| from pathlib import Path | |
| from typing import Any, Mapping, Optional, Sequence, Union | |
| import numpy as np | |
| Source = Union[str, os.PathLike, bytes, bytearray] | |
| IMAGE_EXT = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".webp", ".bmp", ".gif"} | |
| HTML_EXT = {".html", ".htm", ".xhtml"} | |
| TEXT_EXT = {".txt", ".md", ".markdown", ".text", ".csv", ".tsv", ".rst", ".log"} | |
| KINDS = ("pdf", "image", "html", "text") | |
| OCR_ENGINES = ("tesseract", "rapidocr") | |
| class IngestResult: | |
| text: str # the whole document, pages separated by a blank line | |
| pages: list[str] # text per page (one entry for html / text) | |
| method_per_page: list[str] # "text", "text+tables(2)", "ocr:tesseract", "ocr:rapidocr", "html", "plain" | |
| seconds: float # wall time of the whole conversion | |
| warnings: list[str] = field(default_factory=list) | |
| kind: str = "" # "pdf" | "image" | "html" | "text" | |
| ocr_confidence: list[Optional[float]] = field(default_factory=list) # mean word confidence 0-100 per page (None: not OCR'd / unknown) | |
| def to_dict(self, with_text: bool = True) -> dict: | |
| d = asdict(self) | |
| if not with_text: | |
| d.pop("text"), d.pop("pages") | |
| d["chars"] = len(self.text) | |
| d["n_pages"] = len(self.pages) | |
| return d | |
| # ================================================================ public API | |
| def to_text(src: "Source | Sequence[Source]", kind: Optional[str] = None, *, ocr: str = "tesseract", | |
| lang: str = "eng", dpi: int = 300, min_chars: int = 30, tables: bool = True, | |
| preprocess: bool = True, strip_headers: bool = True, max_pages: Optional[int] = None, | |
| ocr_timeout: float = 120.0) -> IngestResult: | |
| """Convert a document to text. | |
| ``src``: a path, raw bytes, or a list of them (the pages of one document, e.g. page scans). | |
| ``kind``: "pdf" | "image" | "html" | "text" (default: from the extension, else sniffed from | |
| the content). ``ocr``: "tesseract" or "rapidocr". ``dpi``: render resolution for PDF pages | |
| that need OCR. ``min_chars``: a PDF page with fewer non-space characters in its text layer is | |
| OCR'd. ``tables``: markdown rows for ruled PDF tables (pdfplumber). ``preprocess``: image | |
| cleanup before OCR (see module doc). ``strip_headers``: drop running headers/footers of | |
| multi-page documents. ``max_pages``: stop after this many pages (with a warning). | |
| """ | |
| if ocr not in OCR_ENGINES: | |
| raise ValueError(f"ocr must be one of {OCR_ENGINES}, got {ocr!r}") | |
| if kind is not None and kind not in KINDS: | |
| raise ValueError(f"kind must be one of {KINDS}, got {kind!r}") | |
| t0 = time.perf_counter() | |
| opts = dict(ocr=ocr, lang=lang, dpi=dpi, min_chars=min_chars, tables=tables, preprocess=preprocess, | |
| max_pages=max_pages, ocr_timeout=ocr_timeout) | |
| srcs = list(src) if isinstance(src, (list, tuple)) else [src] | |
| if not srcs: | |
| raise ValueError("nothing to ingest") | |
| pages: list[str] = [] | |
| methods: list[str] = [] | |
| confs: list[Optional[float]] = [] | |
| warnings: list[str] = [] | |
| kinds: list[str] = [] | |
| for s in srcs: | |
| data, name = _read(s) | |
| k = kind or detect_kind(data, name) | |
| kinds.append(k) | |
| if k == "pdf": | |
| p, m, c, w = _ingest_pdf(data, **opts) | |
| elif k == "image": | |
| p, m, c, w = _ingest_image(data, **opts) | |
| elif k == "html": | |
| p, m, c, w = [html_to_text(_decode(data))], ["html"], [None], [] | |
| else: | |
| p, m, c, w = [normalize_text(_decode(data))], ["plain"], [None], [] | |
| if len(srcs) > 1 and w: | |
| w = [f"{name or 'input'}: {x}" for x in w] | |
| pages += p | |
| methods += m | |
| confs += c | |
| warnings += w | |
| if max_pages and len(pages) >= max_pages: | |
| if len(pages) > max_pages or s is not srcs[-1]: | |
| warnings.append(f"stopped after max_pages={max_pages} pages") | |
| pages, methods, confs = pages[:max_pages], methods[:max_pages], confs[:max_pages] | |
| break | |
| if strip_headers and len(pages) >= 2: | |
| pages, removed = strip_repeated_lines(pages) | |
| if removed: | |
| warnings.append(f"removed {removed} repeated header/footer line(s)") | |
| text = "\n\n".join(p for p in pages if p.strip()) | |
| if not text.strip(): | |
| warnings.append("no text extracted") | |
| ocr_conf = [c for c in confs if c is not None] | |
| if ocr_conf and min(ocr_conf) < 60: | |
| warnings.append(f"low OCR confidence on some pages (min {min(ocr_conf):.0f}/100): check the text") | |
| k = kinds[0] if len(set(kinds)) == 1 else "mixed" | |
| return IngestResult(text, pages, methods, time.perf_counter() - t0, warnings, k, confs) | |
| def predict_file(model, path: "Source | Sequence[Source]", questions: Mapping[str, Any], *, | |
| ingest_kwargs: Optional[dict] = None, **predict_kwargs) -> dict: | |
| """Ingest ``path`` with :func:`to_text` and answer ``questions`` with ``model.predict``. | |
| When the text is longer than the model's state budget (``max_len`` minus the question budget) | |
| and the caller did not choose ``long_state``, it is read in overlapping windows | |
| (``long_state="chunk"``) instead of being cut in the middle. The result is the usual predict | |
| output plus ``"ingest"``: {pages, method_per_page, seconds, warnings, chars, kind, | |
| ocr_confidence, long_state}. | |
| """ | |
| ing = to_text(path, **(ingest_kwargs or {})) | |
| text = ing.text if ing.text.strip() else "(no text could be extracted from the document)" | |
| kw = dict(predict_kwargs) | |
| params = _params(model.predict) | |
| if "long_state" in params and "long_state" not in kw and _too_long(model, text): | |
| kw["long_state"] = "chunk" | |
| res = model.predict(text, questions, **kw) | |
| info = ing.to_dict(with_text=False) | |
| info["long_state"] = kw.get("long_state", getattr(model, "long_state", None)) | |
| res["ingest"] = info | |
| return res | |
| def detect_kind(data: bytes, name: Optional[str] = None) -> str: | |
| """"pdf" | "image" | "html" | "text" from the file extension, else from the leading bytes.""" | |
| ext = Path(name).suffix.lower() if name else "" | |
| if ext == ".pdf": | |
| return "pdf" | |
| if ext in IMAGE_EXT: | |
| return "image" | |
| if ext in HTML_EXT: | |
| return "html" | |
| if ext in TEXT_EXT: | |
| return "text" | |
| head = bytes(data[:1024]) | |
| if head.lstrip().startswith(b"%PDF") or b"%PDF-" in head[:1024]: | |
| return "pdf" | |
| if (head.startswith(b"\x89PNG") or head.startswith(b"\xff\xd8\xff") or head[:4] in (b"II*\x00", b"MM\x00*") | |
| or (head.startswith(b"RIFF") and head[8:12] == b"WEBP") or head.startswith(b"GIF8") | |
| or head.startswith(b"BM")): | |
| return "image" | |
| low = head.lower().lstrip() | |
| if low.startswith((b"<!doctype html", b"<html", b"<?xml")) or b"<html" in low or b"<body" in low: | |
| return "html" | |
| return "text" | |
| # ================================================================ helpers: io | |
| def _read(s: Source) -> tuple[bytes, Optional[str]]: | |
| if isinstance(s, (bytes, bytearray)): | |
| return bytes(s), None | |
| p = Path(s) | |
| return p.read_bytes(), p.name | |
| def _decode(data: bytes) -> str: | |
| if data.startswith(b"\xef\xbb\xbf"): | |
| data = data[3:] | |
| for enc in ("utf-8", "cp1252"): | |
| try: | |
| return data.decode(enc) | |
| except UnicodeDecodeError: | |
| pass | |
| return data.decode("latin-1") | |
| def _params(fn) -> set[str]: | |
| try: | |
| return set(inspect.signature(fn).parameters) | |
| except (TypeError, ValueError): | |
| return set() | |
| def _too_long(model, text: str) -> bool: | |
| ser = getattr(model, "ser", None) | |
| tok, max_len = getattr(ser, "tok", None), getattr(model, "max_len", None) | |
| if tok is None or not max_len: | |
| return False | |
| budget = max_len - getattr(ser, "head_max_len", 0) - 2 | |
| # cheap bound first: ~>=2 chars per token for English | |
| if len(text) < 2 * budget: | |
| return False | |
| return len(tok(text, add_special_tokens=False)["input_ids"]) > budget | |
| # ================================================================ text cleanup | |
| _LIGATURES = {"ff": "ff", "fi": "fi", "fl": "fl", "ffi": "ffi", "ffl": "ffl", "ſt": "st", | |
| "st": "st"} | |
| _INVISIBLE = dict.fromkeys(map(ord, ""), None) | |
| _BULLET = re.compile(r"^\s*(?:[-*•▪●‣⁃◦]|\(?\d{1,3}[.)]|\(?[a-z][.)])\s+") | |
| _ROW = re.compile(r"(^\s*\|)|( \| )") | |
| def _is_row(line: str) -> bool: | |
| return bool(_ROW.search(line)) | |
| def normalize_text(s: str) -> str: | |
| """Whitespace/unicode normalization only (no re-flow): for text and markdown inputs.""" | |
| s = s.replace("\r\n", "\n").replace("\r", "\n").translate(_INVISIBLE) | |
| for k, v in _LIGATURES.items(): | |
| s = s.replace(k, v) | |
| s = s.replace(" ", " ").replace("\t", " ") | |
| lines = [re.sub(r"[ \f\v]+$", "", ln) for ln in s.split("\n")] | |
| s = "\n".join(lines) | |
| s = re.sub(r"\n{3,}", "\n\n", s) | |
| return s.strip() | |
| def clean_text(s: str, reflow: bool = True) -> str: | |
| """Cleanup for extracted / OCR'd text: unicode fixes, de-hyphenation, paragraph re-flow, | |
| removal of lines without any letter or digit (OCR specks), whitespace normalization.""" | |
| s = s.replace("\r\n", "\n").replace("\r", "\n").translate(_INVISIBLE) | |
| for k, v in _LIGATURES.items(): | |
| s = s.replace(k, v) | |
| s = s.replace(" ", " ").replace("\t", " ").replace("\f", "\n") | |
| lines = [] | |
| for ln in s.split("\n"): | |
| ln = re.sub(r" {2,}", " ", ln).strip() | |
| if ln and not re.search(r"[^\W_]", ln): # no letter/digit at all: rule, speck, "|", "~ ." | |
| continue | |
| lines.append(ln) | |
| s = "\n".join(lines) | |
| # de-hyphenate a word split across a line break: "infor-\nmation" -> "information" | |
| s = re.sub(r"([A-Za-z]{2,})[-‐‑]\n[ ]*([a-z]{2,})", r"\1\2", s) | |
| paras = re.split(r"\n\s*\n", s) | |
| out = [] | |
| for p in paras: | |
| pl = [ln for ln in p.split("\n") if ln.strip()] | |
| if not pl: | |
| continue | |
| if not reflow: | |
| out.append("\n".join(pl)) | |
| continue | |
| merged: list[str] = [] | |
| for ln in pl: | |
| keep_break = (not merged or _is_row(ln) or _is_row(merged[-1]) or _BULLET.match(ln) | |
| or ln.startswith("#") or merged[-1].startswith("#")) | |
| if keep_break: | |
| merged.append(ln) | |
| else: | |
| merged[-1] = merged[-1] + " " + ln | |
| out.append("\n".join(merged)) | |
| return "\n\n".join(out).strip() | |
| def _hf_key(line: str) -> str: | |
| """Header/footer comparison key: lowercase letters and '#' for each digit run.""" | |
| return re.sub(r"\d+", "#", re.sub(r"[^a-z0-9]+", "", line.lower())) | |
| def _similar(a: str, b: str) -> bool: | |
| if a == b: | |
| return True | |
| if min(len(a), len(b)) < 6: | |
| return False | |
| import difflib | |
| return difflib.SequenceMatcher(None, a, b, autojunk=False).ratio() >= 0.8 | |
| def strip_repeated_lines(pages: list[str], edge: int = 2, min_frac: float = 0.5) -> tuple[list[str], int]: | |
| """Remove running headers / footers: a line among the first or last ``edge`` lines of a page | |
| that (digits masked, fuzzy-matched so OCR noise does not hide it) occurs at a page edge on | |
| >= ``min_frac`` of the pages (and on >= 2 pages).""" | |
| if len(pages) < 2: | |
| return pages, 0 | |
| split = [p.split("\n") for p in pages] | |
| edges = [] | |
| for lines in split: | |
| idx = [i for i, ln in enumerate(lines) if ln.strip()] | |
| e = sorted(set(idx[:edge] + idx[-edge:])) | |
| # headers/footers are short lines; long ones are (re-flowed) body text | |
| edges.append([(i, _hf_key(lines[i])) for i in e if 0 < len(_hf_key(lines[i])) and len(lines[i]) <= 100]) | |
| need = max(2, int(np.ceil(min_frac * len(pages)))) | |
| drop: list[set[int]] = [set() for _ in pages] | |
| for p, ed in enumerate(edges): | |
| for i, k in ed: | |
| hits = sum(any(_similar(k, k2) for _, k2 in edges[q]) for q in range(len(pages)) if q != p) + 1 | |
| if hits >= need: | |
| drop[p].add(i) | |
| if len(drop[p]) >= sum(1 for ln in split[p] if ln.strip()): | |
| drop[p] = set() # never empty a page (duplicate or near-empty pages) | |
| removed = sum(len(d) for d in drop) | |
| if not removed: | |
| return pages, 0 | |
| out = [] | |
| for lines, d in zip(split, drop): | |
| keep = [ln for i, ln in enumerate(lines) if i not in d] | |
| out.append(re.sub(r"\n{3,}", "\n\n", "\n".join(keep)).strip()) | |
| return out, removed | |
| # ================================================================ layout: lines -> text | |
| class _Line: | |
| x0: float | |
| y0: float | |
| x1: float | |
| y1: float | |
| text: str | |
| par: int = 0 # paragraph id inside its block; a change starts a new paragraph | |
| blk: int = 0 # source block id (set when blocks are zipped) | |
| def h(self) -> float: | |
| return max(1.0, self.y1 - self.y0) | |
| def yc(self) -> float: | |
| return (self.y0 + self.y1) / 2 | |
| def _numeric(text: str) -> bool: | |
| t = re.sub(r"\s", "", text) | |
| return bool(t) and sum(ch.isdigit() or ch in "$%.,()-–—+/:€£" for ch in t) >= 0.6 * len(t) | |
| _NUMBER = re.compile(r"^[(\-\u2013\u2212+$\u20ac\u00a3]*\d[\d.,]*%?\)?[%,;:]?$") | |
| def _cell_breaks(texts: list[str], gaps: list[float], h: float, prose: bool = False) -> list[bool]: | |
| """Which horizontal gaps inside one text line separate table cells (joined with " | "). | |
| Inside a prose block only very wide gaps (> 4 word heights) count, so stretched justified | |
| lines stay prose. Elsewhere a gap must be wide (> 1.5) and then either very wide, next to a | |
| number token, or much wider than the line's narrowest gap.""" | |
| if not gaps: | |
| return [] | |
| gmin = min(gaps) | |
| out = [] | |
| for j, g in enumerate(gaps): | |
| if prose: | |
| out.append(g > 4.0 * h) | |
| continue | |
| a, b = _NUMBER.match(texts[j].strip()), _NUMBER.match(texts[j + 1].strip()) | |
| if a and b: # two numbers side by side: adjacent numeric cells (padding ~1 word height) | |
| out.append(g > 1.0 * h) | |
| continue | |
| out.append(g > 1.5 * h and (g > 4.0 * h or bool(a or b) or (len(gaps) > 1 and g > 3.0 * max(gmin, 0.05 * h)))) | |
| return out | |
| def _nwords(text: str) -> int: | |
| return sum(1 for w in text.split() if w != "|") | |
| def _prose_lines(word_counts: list[int]) -> bool: | |
| """A block of running text: median >= 5 words per line over >= 3 lines.""" | |
| return len(word_counts) >= 3 and statistics.median(word_counts) >= 5 | |
| def _y_overlap(a0, a1, b0, b1) -> float: | |
| return max(0.0, min(a1, b1) - max(a0, b0)) | |
| def _bbox(lines: list[_Line]) -> tuple[float, float, float, float]: | |
| return (min(l.x0 for l in lines), min(l.y0 for l in lines), max(l.x1 for l in lines), max(l.y1 for l in lines)) | |
| def _aligned(a: list[_Line], b: list[_Line]) -> float: | |
| """Fraction of lines of the smaller block with a y-matching line in the other block.""" | |
| small, big = (a, b) if len(a) <= len(b) else (b, a) | |
| hit = sum(any(abs(s.yc - t.yc) < 0.5 * min(s.h, t.h) for t in big) for s in small) | |
| return hit / max(1, len(small)) | |
| def _zip_blocks(blocks: list[list[_Line]], page_w: float) -> list[tuple[list[_Line], bool]]: | |
| """Merge runs of side-by-side blocks whose lines line up (table columns emitted one after | |
| another) into one table block. Returns [(lines, is_table)] in reading order.""" | |
| out: list[tuple[list[_Line], bool]] = [] | |
| group: list[list[_Line]] = [] | |
| def narrow(b): | |
| return max(l.x1 - l.x0 for l in b) < 0.4 * page_w | |
| def prose(b): # columns of running text (>= 5 words per line) are never zipped into rows | |
| return statistics.median(_nwords(l.text) for l in b) >= 5 | |
| def flush(): | |
| if not group: | |
| return | |
| if len(group) == 1: | |
| out.append((group[0], False)) | |
| else: | |
| lines = [] | |
| for i, b in enumerate(group): | |
| for l in b: | |
| l.blk = i | |
| lines.append(l) | |
| out.append((lines, True)) | |
| group.clear() | |
| for b in blocks: | |
| if group: | |
| prev = group[-1] | |
| gx0, gy0, gx1, gy1 = _bbox([l for g in group for l in g]) | |
| bx0, by0, bx1, by1 = _bbox(b) | |
| px0, py0, px1, py1 = _bbox(prev) | |
| hmin = min(statistics.median(l.h for l in b), statistics.median(l.h for l in prev)) | |
| right = bx0 >= px1 - 0.5 * hmin | |
| ov = _y_overlap(gy0, gy1, by0, by1) / max(1.0, min(gy1 - gy0, by1 - by0)) | |
| both_single = len(b) == 1 and len(prev) == 1 | |
| numeric = sum(_numeric(l.text) for l in b) >= 0.6 * len(b) | |
| if right and ov >= 0.5 and _aligned(prev, b) >= 0.6 and not (prose(b) or prose(prev)) and ( | |
| both_single or numeric or (narrow(b) and narrow(prev))): | |
| group.append(b) | |
| continue | |
| flush() | |
| group.append(b) | |
| flush() | |
| return out | |
| def _join_row(cells: list[_Line], table: bool, prose: bool = False) -> str: | |
| cells = sorted(cells, key=lambda l: l.x0) | |
| h = min(l.h for l in cells) | |
| gaps = [b.x0 - a.x1 for a, b in zip(cells, cells[1:])] | |
| # separate line boxes on one row: cells of a zipped table (different source blocks), or a | |
| # wide gap by the rule of _cell_breaks (boxes are measured in box heights, ~1.4 word heights) | |
| breaks = _cell_breaks([c.text for c in cells], gaps, 0.7 * h, prose=prose and not table) | |
| txt = cells[0].text | |
| for a, b, br in zip(cells, cells[1:], breaks): | |
| txt += (" | " if (table and a.blk != b.blk) or br else " ") + b.text | |
| return txt | |
| def _rows(lines: list[_Line]) -> list[str]: | |
| """Cluster lines by vertical position into rows (for zipped table blocks).""" | |
| hmed = statistics.median(l.h for l in lines) | |
| rows: list[list[_Line]] = [] | |
| for l in sorted(lines, key=lambda l: l.yc): | |
| if rows and abs(l.yc - statistics.mean(x.yc for x in rows[-1])) < 0.5 * hmed: | |
| rows[-1].append(l) | |
| else: | |
| rows.append([l]) | |
| return [_join_row(r, True) for r in rows] | |
| def _seq_text(lines: list[_Line]) -> str: | |
| """Lines in reading order; a line that continues the previous one on the same baseline (a | |
| table cell to its right) is appended to it; paragraph changes give a blank line.""" | |
| out: list[str] = [] | |
| cur: list[_Line] = [] | |
| prose = _prose_lines([_nwords(l.text) for l in lines]) | |
| for l in lines: | |
| if cur: | |
| last = cur[-1] | |
| same_row = (l.x0 >= last.x1 - 0.2 * l.h and | |
| _y_overlap(last.y0, last.y1, l.y0, l.y1) >= 0.5 * min(last.h, l.h)) | |
| if same_row: | |
| cur.append(l) | |
| continue | |
| out.append(_join_row(cur, False, prose)) | |
| if l.par != last.par: | |
| out.append("") | |
| cur = [l] | |
| if cur: | |
| out.append(_join_row(cur, False, prose)) | |
| return "\n".join(out) | |
| def layout_text(blocks: list[list[_Line]], page_w: float) -> str: | |
| parts: list[str] = [] | |
| for lines, is_table in _zip_blocks([b for b in blocks if b], page_w): | |
| p = "\n".join(_rows(lines)) if is_table else _seq_text(lines) | |
| if not p.strip(): | |
| continue | |
| if parts and _is_row(parts[-1].rsplit("\n", 1)[-1]) and _is_row(p.split("\n", 1)[0]): | |
| parts[-1] += "\n" + p # consecutive table rows stay one block | |
| else: | |
| parts.append(p) | |
| return "\n\n".join(parts) | |
| def markdown_table(rows: list[list[Optional[str]]]) -> str: | |
| rows = [[re.sub(r"\s+", " ", c or "").strip() for c in r] for r in rows] | |
| rows = [r for r in rows if any(r)] | |
| if not rows: | |
| return "" | |
| n = max(len(r) for r in rows) | |
| rows = [r + [""] * (n - len(r)) for r in rows] | |
| # drop columns that are empty everywhere (merged-cell artefacts) | |
| keep = [j for j in range(n) if any(r[j] for r in rows)] | |
| rows = [[r[j] for j in keep] for r in rows] | |
| # no "| --- |" separator row: it carries no information for the model | |
| return "\n".join("| " + " | ".join(r) + " |" for r in rows) | |
| def _ingest_pdf(data: bytes, *, ocr, lang, dpi, min_chars, tables, preprocess, max_pages, ocr_timeout): | |
| import fitz # PyMuPDF | |
| warnings: list[str] = [] | |
| try: | |
| doc = fitz.open(stream=data, filetype="pdf") | |
| except Exception as e: | |
| raise ValueError(f"cannot open PDF: {type(e).__name__}: {e}") from None | |
| if doc.needs_pass and not doc.authenticate(""): | |
| raise ValueError("PDF is password-protected") | |
| plumber = None | |
| if tables: | |
| try: | |
| import pdfplumber | |
| plumber = pdfplumber.open(io.BytesIO(data)) | |
| except Exception as e: # pragma: no cover - broken/unsupported file for pdfplumber | |
| warnings.append(f"table detection unavailable: {type(e).__name__}") | |
| pages, methods, confs = [], [], [] | |
| n = doc.page_count if not max_pages else min(doc.page_count, max_pages) | |
| try: | |
| for i in range(n): | |
| page = doc[i] | |
| ppage = plumber.pages[i] if plumber is not None and i < len(plumber.pages) else None | |
| text, n_tab = _pdf_page_text(page, ppage, warnings) | |
| reason = _needs_ocr(page, text, min_chars) | |
| if reason: | |
| img = _render(page, dpi) | |
| otext, conf = _ocr_image(img, ocr=ocr, lang=lang, dpi=dpi, preprocess=preprocess, | |
| timeout=ocr_timeout) | |
| if len(otext.strip()) >= len(text.strip()): | |
| pages.append(otext) | |
| methods.append(f"ocr:{ocr}") | |
| confs.append(conf) | |
| continue | |
| pages.append(text) | |
| methods.append(f"text+tables({n_tab})" if n_tab else "text") | |
| confs.append(None) | |
| finally: | |
| if plumber is not None: | |
| plumber.close() | |
| if max_pages and doc.page_count > max_pages: | |
| warnings.append(f"PDF has {doc.page_count} pages; read the first {max_pages}") | |
| return pages, methods, confs, warnings | |
| def _needs_ocr(page, text: str, min_chars: int) -> Optional[str]: | |
| body = re.sub(r"\s", "", text) | |
| if len(body) < min_chars: | |
| return "little text" | |
| bad = body.count("�") + len(re.findall(r"\(cid:\d+\)", text)) * 6 | |
| if bad > 0.2 * len(body): | |
| return "unmapped glyphs" | |
| if len(body) < 300: | |
| try: | |
| area = abs(page.rect) | |
| img_area = sum(abs(r) for info in page.get_images(full=True) for r in page.get_image_rects(info[0])) | |
| if area and img_area > 0.5 * area: | |
| return "image page with a thin text layer" | |
| except Exception: | |
| pass | |
| return None | |
| def _render(page, dpi: int) -> np.ndarray: | |
| import fitz | |
| pix = page.get_pixmap(dpi=dpi, colorspace=fitz.csGRAY, alpha=False) | |
| a = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.stride)[:, : pix.width] | |
| return a.copy() | |
| def _pdf_page_text(page, ppage, warnings: list[str]) -> tuple[str, int]: | |
| tabs: list[tuple[tuple[float, float, float, float], str]] = [] | |
| if ppage is not None and page.rotation == 0: | |
| try: | |
| for t in ppage.find_tables(): | |
| rows = t.extract() | |
| if len(rows) >= 2 and max(len(r) for r in rows) >= 2: | |
| md = markdown_table(rows) | |
| if md: | |
| tabs.append((t.bbox, md)) | |
| except Exception as e: # pragma: no cover | |
| warnings.append(f"table detection failed on page {page.number + 1}: {type(e).__name__}") | |
| d = page.get_text("dict", flags=0) # no images; content-stream order | |
| blocks: list[list[_Line]] = [] | |
| placed: set[int] = set() | |
| items: list[Any] = [] # list of block line-lists or ("table", md) | |
| for b in d.get("blocks", []): | |
| if b.get("type", 0) != 0: | |
| continue | |
| cur: list[_Line] = [] | |
| blk_prose = _prose_lines([len("".join(sp.get("text", "") for sp in li.get("spans", [])).split()) | |
| for li in b.get("lines", [])]) | |
| for li in b.get("lines", []): | |
| spans = [s for s in li.get("spans", []) if s.get("text", "").strip()] | |
| if not spans: | |
| continue | |
| x0, y0, x1, y1 = li["bbox"] | |
| cx, cy = (x0 + x1) / 2, (y0 + y1) / 2 | |
| ti = next((k for k, (bb, _) in enumerate(tabs) if bb[0] - 1 <= cx <= bb[2] + 1 and bb[1] - 1 <= cy <= bb[3] + 1), None) | |
| if ti is not None: | |
| if cur: | |
| items.append(cur) | |
| cur = [] | |
| if ti not in placed: | |
| placed.add(ti) | |
| items.append(("table", tabs[ti][1])) | |
| continue | |
| txt = spans[0]["text"] | |
| sgaps = [s2["bbox"][0] - s1["bbox"][2] for s1, s2 in zip(spans, spans[1:])] | |
| size = max(1.0, statistics.median(sp.get("size", 10) for sp in spans)) | |
| breaks = _cell_breaks([sp["text"] for sp in spans], sgaps, 0.7 * size, prose=blk_prose) | |
| for sp, g, br in zip(spans[1:], sgaps, breaks): | |
| sep = " | " if br else ("" if g < 0.15 * size else " ") | |
| txt += sep + sp["text"] | |
| cur.append(_Line(x0, y0, x1, y1, re.sub(r"\s+", " ", txt).strip())) | |
| if cur: | |
| items.append(cur) | |
| for ti, (_, md) in enumerate(tabs): | |
| if ti not in placed: | |
| items.append(("table", md)) | |
| # layout: consecutive line blocks go through layout_text; tables are inserted verbatim | |
| parts: list[str] = [] | |
| run: list[list[_Line]] = [] | |
| for it in items: | |
| if isinstance(it, tuple): | |
| if run: | |
| parts.append(layout_text(run, page.rect.width)) | |
| run = [] | |
| parts.append(it[1]) | |
| else: | |
| run.append(it) | |
| if run: | |
| parts.append(layout_text(run, page.rect.width)) | |
| return clean_text("\n\n".join(parts)), len(tabs) | |
| # ================================================================ images + OCR | |
| def _ingest_image(data: bytes, *, ocr, lang, dpi, min_chars, tables, preprocess, max_pages, ocr_timeout): | |
| from PIL import Image, ImageOps | |
| Image.MAX_IMAGE_PIXELS = 200_000_000 | |
| warnings: list[str] = [] | |
| try: | |
| im = Image.open(io.BytesIO(data)) | |
| except Exception as e: | |
| raise ValueError(f"cannot open image: {type(e).__name__}: {e}") from None | |
| n_frames = getattr(im, "n_frames", 1) | |
| pages, methods, confs = [], [], [] | |
| for f in range(n_frames if not max_pages else min(n_frames, max_pages)): | |
| if n_frames > 1: | |
| im.seek(f) | |
| fr = ImageOps.exif_transpose(im) if f == 0 else im | |
| src_dpi = None | |
| try: | |
| src_dpi = float(fr.info.get("dpi", (0, 0))[0]) or None | |
| except Exception: | |
| pass | |
| gray = np.asarray(fr.convert("L")) | |
| text, conf = _ocr_image(gray, ocr=ocr, lang=lang, dpi=src_dpi, preprocess=preprocess, timeout=ocr_timeout) | |
| pages.append(text) | |
| methods.append(f"ocr:{ocr}") | |
| confs.append(conf) | |
| return pages, methods, confs, warnings | |
| def _ocr_image(gray: np.ndarray, *, ocr: str, lang: str, dpi: Optional[float], preprocess: bool, | |
| timeout: float) -> tuple[str, Optional[float]]: | |
| if preprocess: | |
| gray, scale = prepare_image(gray, upscale=ocr == "tesseract") # RapidOCR rescales internally | |
| dpi = (dpi or 300) * scale if dpi else None | |
| if ocr == "rapidocr": | |
| blocks, conf = _rapidocr_blocks(gray) | |
| else: | |
| blocks, conf = _tesseract_blocks(gray, lang=lang, dpi=dpi, timeout=timeout) | |
| return clean_text(layout_text(blocks, gray.shape[1])), conf | |
| # ---------------------------------------------------------------- preprocessing | |
| def prepare_image(gray: np.ndarray, upscale: bool = True) -> tuple[np.ndarray, float]: | |
| """Photo/scan cleanup before OCR: perspective correction, denoising, illumination flattening, | |
| deskew and (``upscale``) enlarging small text. Returns (image, scale factor applied).""" | |
| import cv2 | |
| g = gray | |
| g = _dewarp_page(g) | |
| g = _denoise(g) | |
| g = _flatten_illumination(g) | |
| g = _deskew(g) | |
| xh = _text_height(g) if upscale else None | |
| scale = 1.0 | |
| if xh and xh < 20: | |
| scale = min(3.0, 26.0 / xh, 6000.0 / max(g.shape)) | |
| if scale > 1.15: | |
| g = cv2.resize(g, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC) | |
| else: | |
| scale = 1.0 | |
| return g, scale | |
| def _small(g: np.ndarray, target: int = 1000) -> tuple[np.ndarray, float]: | |
| import cv2 | |
| f = target / max(g.shape) | |
| if f >= 1: | |
| return g, 1.0 | |
| return cv2.resize(g, None, fx=f, fy=f, interpolation=cv2.INTER_AREA), f | |
| def _quad_angles_ok(q: np.ndarray, tol: float = 25.0) -> bool: | |
| """Every interior angle of the quadrilateral within 90 +- tol degrees (a page seen in | |
| perspective; a region cut off by a shadow edge or a fold is rejected).""" | |
| for i in range(4): | |
| a, b, c = q[i - 1], q[i], q[(i + 1) % 4] | |
| v1, v2 = a - b, c - b | |
| cos = float(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-9)) | |
| if abs(np.degrees(np.arccos(np.clip(cos, -1, 1))) - 90) > tol: | |
| return False | |
| return True | |
| def _order_quad(q: np.ndarray) -> np.ndarray: | |
| s, d = q.sum(1), np.diff(q, axis=1).ravel() | |
| return np.array([q[np.argmin(s)], q[np.argmin(d)], q[np.argmax(s)], q[np.argmax(d)]], np.float32) | |
| def _page_quad(s: np.ndarray) -> Optional[np.ndarray]: | |
| """Largest plausible page quadrilateral in a (downscaled) photo, from several binarizations: | |
| page-vs-background edges (robust to shadows falling across the page) and global thresholds.""" | |
| import cv2 | |
| H, W = s.shape | |
| blur = cv2.GaussianBlur(s, (5, 5), 0) | |
| masks = [] | |
| edges = cv2.Canny(blur, 30, 90) | |
| masks.append(cv2.dilate(edges, np.ones((5, 5), np.uint8))) | |
| otsu, _ = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) | |
| for t in (otsu, 0.75 * otsu, 0.55 * otsu): | |
| _, th = cv2.threshold(blur, t, 255, cv2.THRESH_BINARY) | |
| masks.append(cv2.morphologyEx(th, cv2.MORPH_CLOSE, np.ones((9, 9), np.uint8))) | |
| border = 0.006 * max(H, W) | |
| for m in masks: # priority order: edges first, then thresholds from strict to loose | |
| cnts, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| for c in sorted(cnts, key=cv2.contourArea, reverse=True)[:2]: | |
| hull = cv2.convexHull(c) | |
| area = cv2.contourArea(hull) | |
| if area < 0.25 * H * W or area > 0.97 * H * W: | |
| continue | |
| peri = cv2.arcLength(hull, True) | |
| for eps in (0.02, 0.03, 0.05): | |
| ap = cv2.approxPolyDP(hull, eps * peri, True) | |
| if len(ap) == 4: | |
| q = _order_quad(ap.reshape(4, 2).astype(np.float32)) | |
| inside = bool(((q[:, 0] > border) & (q[:, 0] < W - 1 - border) & | |
| (q[:, 1] > border) & (q[:, 1] < H - 1 - border)).all()) | |
| if inside and _quad_angles_ok(q) and cv2.contourArea(q) >= 0.9 * area: | |
| return q # a corner on the frame border = leaked into the background | |
| break | |
| return None | |
| def _dewarp_page(g: np.ndarray) -> np.ndarray: | |
| """Perspective-correct a photographed page: the largest plausible page quadrilateral | |
| (>= 25% of the frame, near-right angles) that is clearly brighter than its surroundings.""" | |
| import cv2 | |
| s, f = _small(g) | |
| quad = _page_quad(s) | |
| if quad is None: | |
| return g | |
| mask = np.zeros_like(s) | |
| cv2.fillConvexPoly(mask, quad.astype(np.int32), 255) | |
| inside, outside = s[mask > 0], s[mask == 0] | |
| if outside.size < 0.03 * s.size or float(np.median(inside)) - float(np.median(outside)) < 40: | |
| return g | |
| tl, tr, br, bl = quad / f | |
| w = int(max(np.linalg.norm(tr - tl), np.linalg.norm(br - bl))) | |
| h = int(max(np.linalg.norm(bl - tl), np.linalg.norm(br - tr))) | |
| if w < 100 or h < 100: | |
| return g | |
| M = cv2.getPerspectiveTransform(np.array([tl, tr, br, bl], np.float32), | |
| np.array([[0, 0], [w - 1, 0], [w - 1, h - 1], [0, h - 1]], np.float32)) | |
| return cv2.warpPerspective(g, M, (w, h), flags=cv2.INTER_CUBIC, borderValue=255) | |
| def _denoise(g: np.ndarray) -> np.ndarray: | |
| """3x3 median filter when the paper shows pixel noise / specks (noisy scans, grainy photos).""" | |
| import cv2 | |
| med = cv2.medianBlur(g, 3) | |
| thr, _ = cv2.threshold(med, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) | |
| bg = med > thr + 10 | |
| if bg.mean() < 0.2: | |
| return g | |
| noise = float(np.std(g[bg].astype(np.float32) - med[bg])) | |
| return med if noise > 3.0 else g | |
| def _flatten_illumination(g: np.ndarray) -> np.ndarray: | |
| """Divide by a smooth background estimate when the paper brightness varies across the page.""" | |
| import cv2 | |
| s, f = _small(g, 600) | |
| bg = cv2.medianBlur(cv2.morphologyEx(s, cv2.MORPH_CLOSE, np.ones((15, 15), np.uint8)), 31) | |
| lo, hi = np.percentile(bg, [5, 95]) | |
| if hi - lo < 30: | |
| return g | |
| bg_full = cv2.resize(bg, (g.shape[1], g.shape[0]), interpolation=cv2.INTER_LINEAR).astype(np.float32) | |
| out = np.clip(g.astype(np.float32) / np.maximum(bg_full, 1) * 245.0, 0, 255) | |
| return out.astype(np.uint8) | |
| def _deskew(g: np.ndarray, max_angle: float = 5.0) -> np.ndarray: | |
| """Projection-profile deskew: the angle that maximizes the variance of row ink sums.""" | |
| import cv2 | |
| s, _ = _small(g, 1200) | |
| _, th = cv2.threshold(s, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) | |
| if th.mean() < 0.5: # (almost) blank | |
| return g | |
| h, w = th.shape | |
| center = (w / 2, h / 2) | |
| def score(a): | |
| M = cv2.getRotationMatrix2D(center, a, 1.0) | |
| r = cv2.warpAffine(th, M, (w, h), flags=cv2.INTER_NEAREST, borderValue=0) | |
| return float(np.var(r.sum(1, dtype=np.float64))) | |
| angles = np.arange(-max_angle, max_angle + 1e-6, 0.5) | |
| best = max(angles, key=score) | |
| fine = np.arange(best - 0.4, best + 0.41, 0.1) | |
| best = float(max(fine, key=score)) | |
| if abs(best) < 0.25: | |
| return g | |
| M = cv2.getRotationMatrix2D((g.shape[1] / 2, g.shape[0] / 2), best, 1.0) | |
| return cv2.warpAffine(g, M, (g.shape[1], g.shape[0]), flags=cv2.INTER_CUBIC, borderValue=255) | |
| def _text_height(g: np.ndarray) -> Optional[float]: | |
| """Median height of character-sized connected components (~ x-height), in pixels.""" | |
| import cv2 | |
| _, th = cv2.threshold(g, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) | |
| n, _, stats, _ = cv2.connectedComponentsWithStats(th, connectivity=8) | |
| if n < 20: | |
| return None | |
| hs = stats[1:, cv2.CC_STAT_HEIGHT] | |
| ws = stats[1:, cv2.CC_STAT_WIDTH] | |
| ok = (hs >= 3) & (hs <= 150) & (ws <= 3 * hs) & (ws >= 1) | |
| if ok.sum() < 20: | |
| return None | |
| return float(np.median(hs[ok])) | |
| # ---------------------------------------------------------------- engines | |
| def _tesseract_blocks(gray: np.ndarray, lang: str = "eng", dpi: Optional[float] = None, psm: int = 3, | |
| timeout: float = 120.0) -> tuple[list[list[_Line]], Optional[float]]: | |
| import cv2 | |
| exe = shutil.which("tesseract") | |
| if not exe: | |
| raise RuntimeError("tesseract is not installed (apt install tesseract-ocr), or use ocr='rapidocr'") | |
| ok, png = cv2.imencode(".png", gray) | |
| env = dict(os.environ) | |
| env.setdefault("OMP_THREAD_LIMIT", "1") | |
| cmd = [exe, "stdin", "stdout", "-l", lang, "--psm", str(psm)] | |
| if dpi: | |
| cmd += ["--dpi", str(int(round(dpi)))] | |
| else: | |
| cmd += ["--dpi", "300"] | |
| cmd.append("tsv") | |
| r = subprocess.run(cmd, input=png.tobytes(), capture_output=True, timeout=timeout, env=env) | |
| if r.returncode != 0: | |
| raise RuntimeError(f"tesseract failed: {r.stderr.decode(errors='replace')[:300]}") | |
| rows = r.stdout.decode("utf-8", errors="replace").splitlines()[1:] | |
| lines: dict[tuple[int, int, int], list[tuple[int, int, int, int, float, str]]] = {} | |
| order: list[tuple[int, int, int]] = [] | |
| confs, weights = [], [] | |
| for row in rows: | |
| f = row.split("\t") | |
| if len(f) < 12 or f[0] != "5": | |
| continue | |
| txt = f[11].strip() | |
| conf = float(f[10]) | |
| if not txt or conf < 0: | |
| continue | |
| key = (int(f[2]), int(f[3]), int(f[4])) | |
| if key not in lines: | |
| lines[key] = [] | |
| order.append(key) | |
| x, y, w, h = int(f[6]), int(f[7]), int(f[8]), int(f[9]) | |
| lines[key].append((x, y, x + w, y + h, conf, txt)) | |
| confs.append(conf) | |
| weights.append(len(txt)) | |
| blocks: dict[int, list[_Line]] = {} | |
| border: list[int] = [] | |
| wc: dict[int, list[int]] = {} | |
| for key in order: | |
| wc.setdefault(key[0], []).append(len(lines[key])) | |
| for key in order: | |
| ws = _fix_pipes(sorted(lines[key], key=lambda w: w[0])) | |
| if not ws: | |
| continue | |
| hs = sorted(w[3] - w[1] for w in ws) | |
| hmed = hs[len(hs) // 2] | |
| txt = ws[0][5] | |
| breaks = _cell_breaks([w[5] for w in ws], [b[0] - a[2] for a, b in zip(ws, ws[1:])], hmed, | |
| prose=_prose_lines(wc[key[0]])) | |
| for b, br in zip(ws[1:], breaks): | |
| txt += (" | " if br else " ") + b[5] | |
| ln = _Line(min(w[0] for w in ws), min(w[1] for w in ws), max(w[2] for w in ws), max(w[3] for w in ws), | |
| txt, par=key[1]) | |
| if key[0] not in blocks: | |
| blocks[key[0]] = [] | |
| border.append(key[0]) | |
| blocks[key[0]].append(ln) | |
| conf = float(np.average(confs, weights=weights)) if confs else None | |
| return [blocks[b] for b in border], conf | |
| def _fix_pipes(ws: list[tuple]) -> list[tuple]: | |
| """Tesseract reads a capital I as "|": a lone "|" followed at a normal word gap by a lowercase | |
| word becomes "I"; other lone bars (table rules) are dropped (cells are split by gaps anyway).""" | |
| out = [] | |
| for j, w in enumerate(ws): | |
| if w[5] in ("|", "||", "|,"): | |
| nxt = ws[j + 1] if j + 1 < len(ws) else None | |
| h = w[3] - w[1] | |
| if nxt is not None and nxt[5][:1].islower() and nxt[0] - w[2] < 1.2 * h: | |
| out.append(w[:5] + ("I" + w[5][1:].replace("|", ""),)) | |
| continue | |
| out.append(w) | |
| return out | |
| _RAPID = None | |
| def _rapidocr_engine(): | |
| global _RAPID | |
| if _RAPID is None: | |
| try: | |
| from rapidocr_onnxruntime import RapidOCR | |
| except ImportError as e: # pragma: no cover | |
| raise RuntimeError("ocr='rapidocr' needs `pip install rapidocr_onnxruntime`") from e | |
| n = int(os.environ.get("JEVLIKE_OCR_THREADS", "0") or 0) | |
| # no dilation of the detector's text map: it merges adjacent lines of blurred text (photos) | |
| kw: dict[str, Any] = {"det_use_dilation": False} | |
| if n > 0: | |
| kw.update(intra_op_num_threads=n, inter_op_num_threads=1) | |
| rec = os.environ.get("JEVLIKE_RAPIDOCR_REC") | |
| if rec: | |
| kw["rec_model_path"] = rec | |
| _RAPID = RapidOCR(**kw) | |
| return _RAPID | |
| def _rapidocr_blocks(gray: np.ndarray) -> tuple[list[list[_Line]], Optional[float]]: | |
| import cv2 | |
| eng = _rapidocr_engine() | |
| # no 180-degree line classifier: pages are upright after deskew, and the classifier flips | |
| # (and then garbles or drops) some long text lines | |
| res, _ = eng(cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR), use_cls=False) | |
| lines = [] | |
| confs, weights = [], [] | |
| for box, txt, score in res or []: | |
| txt = str(txt).strip() | |
| if not txt: | |
| continue | |
| pts = np.asarray(box, dtype=float) | |
| lines.append(_Line(pts[:, 0].min(), pts[:, 1].min(), pts[:, 0].max(), pts[:, 1].max(), txt)) | |
| confs.append(float(score) * 100) | |
| weights.append(len(txt)) | |
| conf = float(np.average(confs, weights=weights)) if confs else None | |
| if not lines: | |
| return [], conf | |
| return _xy_cut(lines), conf | |
| def _gaps(iv: list[tuple[float, float]]) -> list[tuple[float, float]]: | |
| """Whitespace gaps between the union of 1-D intervals: [(width, midpoint)].""" | |
| iv = sorted(iv) | |
| out, end = [], iv[0][1] | |
| for a, b in iv[1:]: | |
| if a > end: | |
| out.append((a - end, (end + a) / 2)) | |
| end = max(end, b) | |
| return out | |
| def _interior(ls: list[_Line], a: float, b: float) -> bool: | |
| """A band [a, b] crossed by some boxes is a real gutter only if every crossing box has | |
| column text above and below it on both sides (a merged line inside two columns), not a | |
| full-width line at the top or bottom of the region (a title or a paragraph above a table: | |
| those are split off by a horizontal cut instead).""" | |
| cross = [l for l in ls if l.x0 < a and l.x1 > b] | |
| left = [l for l in ls if l.x1 <= a] | |
| right = [l for l in ls if l.x0 >= b] | |
| if not left or not right: | |
| return False | |
| for c in cross: | |
| for side in (left, right): | |
| if not (any(l.yc < c.y0 for l in side) and any(l.yc > c.y1 for l in side)): | |
| return False | |
| return True | |
| def _xy_cut(lines: list[_Line], depth: int = 0) -> list[list[_Line]]: | |
| """Reading order for OCR line boxes without layout info (RapidOCR). Recursive XY-cut: | |
| split into columns at the widest vertical whitespace gap (>= 1 line height); otherwise split | |
| horizontally at the widest gap that lets one side split into columns (a full-width title | |
| above a two-column body, a paragraph above a table). Leaves are blocks read top-to-bottom, | |
| with a new paragraph at vertical gaps > 0.6 line height.""" | |
| h = statistics.median(l.h for l in lines) | |
| def xgap(ls): | |
| """Widest vertical whitespace band; a band crossed by a few boxes (text merged across a | |
| narrow gutter by the detector) still counts when the rest of the region splits there.""" | |
| if len(ls) < 2: | |
| return (0.0, None) | |
| g = _gaps([(l.x0, l.x1) for l in ls]) | |
| best = max(g) if g else (0.0, None) | |
| tol = int(0.03 * len(ls)) | |
| if tol >= 1: | |
| x0, x1 = int(min(l.x0 for l in ls)), int(max(l.x1 for l in ls)) + 1 | |
| cov = np.zeros(x1 - x0 + 1, np.int32) | |
| for l in ls: | |
| cov[int(l.x0) - x0:int(l.x1) - x0 + 1] += 1 | |
| low = cov <= tol | |
| run, start = 0, 0 | |
| for i, v in enumerate(np.append(low, False)): | |
| if v: | |
| if run == 0: | |
| start = i | |
| run += 1 | |
| elif run: | |
| if start > 0 and i < len(low) and run > best[0] and _interior(ls, x0 + start, x0 + i): | |
| best = (float(run), x0 + start + run / 2) | |
| run = 0 | |
| return best | |
| def split(ls): | |
| """Column split of a region: a gutter >= 1 line height with running text (>= 4 words per | |
| box) on both sides. Table columns are not split: their rows are read across.""" | |
| gx, px = xgap(ls) | |
| if gx < 1.0 * h or px is None: | |
| return None | |
| left = [l for l in ls if (l.x0 + l.x1) / 2 < px] | |
| right = [l for l in ls if (l.x0 + l.x1) / 2 >= px] | |
| if not left or not right: | |
| return None | |
| if min(statistics.median(_nwords(l.text) for l in side) for side in (left, right)) < 4: | |
| return None | |
| return left, right | |
| if len(lines) > 1 and depth < 60: | |
| cols = split(lines) | |
| if cols: | |
| return _xy_cut(cols[0], depth + 1) + _xy_cut(cols[1], depth + 1) | |
| for gy, py in sorted(_gaps([(l.y0, l.y1) for l in lines]), reverse=True): | |
| if gy < 0.1 * h: | |
| break | |
| top = [l for l in lines if l.yc < py] | |
| bot = [l for l in lines if l.yc >= py] | |
| if top and bot and (split(top) or split(bot)): | |
| return _xy_cut(top, depth + 1) + _xy_cut(bot, depth + 1) | |
| # OCR cell boxes on the same printed row can differ by a few pixels in y. | |
| # Cluster their baselines before sorting horizontally, or a number can be | |
| # emitted before its row label and break the table association. | |
| row_groups: list[list[_Line]] = [] | |
| for line in sorted(lines, key=lambda l: l.yc): | |
| if row_groups and abs(line.yc - statistics.mean(x.yc for x in row_groups[-1])) < 0.5 * h: | |
| row_groups[-1].append(line) | |
| else: | |
| row_groups.append([line]) | |
| leaf = [line for row in row_groups for line in sorted(row, key=lambda l: l.x0)] | |
| par, prev = 0, None | |
| for l in leaf: | |
| if prev is not None and l.y0 > prev.yc and ( | |
| l.y0 - prev.y1 > 0.5 * min(l.h, prev.h) or max(l.h, prev.h) > 1.3 * min(l.h, prev.h)): | |
| par += 1 | |
| l.par = par | |
| if prev is None or l.y1 > prev.y1: | |
| prev = l | |
| return [leaf] | |
| # ================================================================ HTML | |
| _DROP_TAGS = ["script", "style", "noscript", "iframe", "svg", "canvas", "form", "button", "nav", "aside", | |
| "template", "select", "input", "object", "embed", "dialog", "menu"] | |
| _BOILER = re.compile(r"(^|[-_\s])(nav|navbar|menu|footer|sidebar|side-bar|cookie|consent|subscribe|newsletter|" | |
| r"share|sharing|social|breadcrumb|advert|ads?|promo|related|recommend|comments?|popup|modal|" | |
| r"banner|skip-link|signup|paywall)([-_\s]|$)", re.I) | |
| _BLOCK_TAGS = ["h1", "h2", "h3", "h4", "h5", "h6", "p", "li", "table", "pre", "blockquote", "dt", "dd", | |
| "figcaption", "caption", "address"] | |
| def html_to_text(html: str) -> str: | |
| """Main content of an HTML page as markdown-ish text (headings, paragraphs, lists, tables).""" | |
| from bs4 import BeautifulSoup | |
| try: | |
| soup = BeautifulSoup(html, "lxml") | |
| except Exception: | |
| soup = BeautifulSoup(html, "html.parser") | |
| title = soup.title.get_text(" ", strip=True) if soup.title else "" | |
| for t in soup(_DROP_TAGS): | |
| t.decompose() | |
| for t in soup.find_all(["header", "footer"]): | |
| if not t.find_parent(["article", "main"]): | |
| t.decompose() | |
| for t in soup.find_all(True): | |
| if t.attrs is None or t.name in ("html", "body", "article", "main"): | |
| continue | |
| ident = " ".join([t.get("id") or ""] + list(t.get("class") or [])) | |
| role = (t.get("role") or "").lower() | |
| hidden = t.get("aria-hidden") == "true" or "display:none" in (t.get("style") or "").replace(" ", "") | |
| if hidden or role in ("navigation", "banner", "contentinfo", "complementary", "search") or ( | |
| ident.strip() and _BOILER.search(ident)): | |
| t.decompose() | |
| arts = soup.find_all("article") | |
| root = (max(arts, key=lambda a: len(a.get_text())) if arts else None) or soup.find("main") or soup.body or soup | |
| out: list[str] = [] | |
| for el in root.find_all(_BLOCK_TAGS): | |
| if el.find_parent(_BLOCK_TAGS) is not None: # its text is emitted with the enclosing block | |
| continue | |
| if el.name == "table": | |
| rows = [[c.get_text(" ", strip=True) for c in tr.find_all(["th", "td"])] for tr in el.find_all("tr")] | |
| md = markdown_table(rows) | |
| if md: | |
| out.append(md) | |
| continue | |
| txt = el.get_text(" ", strip=True) | |
| txt = re.sub(r"\s+", " ", txt).strip() | |
| if not txt: | |
| continue | |
| if el.name in ("h1", "h2", "h3", "h4", "h5", "h6"): | |
| out.append("#" * int(el.name[1]) + " " + txt) | |
| elif el.name == "li": | |
| out.append("- " + txt) | |
| elif el.name == "pre": | |
| out.append(el.get_text().strip("\n")) | |
| else: | |
| out.append(txt) | |
| body = "\n\n".join(out) | |
| full = re.sub(r"\s+", " ", root.get_text(" ", strip=True)) | |
| if len(re.sub(r"\s", "", body)) < 0.3 * len(re.sub(r"\s", "", full)): | |
| body = root.get_text("\n", strip=True) # div soup without block tags | |
| if title and not re.search(r"^#\s", body, re.M) and title not in body[:500]: | |
| body = "# " + title + "\n\n" + body | |
| # join consecutive list items into one block | |
| body = re.sub(r"(^- .*)\n\n(?=- )", r"\1\n", body, flags=re.M) | |
| body = re.sub(r"(^- .*)\n\n(?=- )", r"\1\n", body, flags=re.M) | |
| return normalize_text(body) | |
| # ================================================================ questions from JSON + CLI | |
| def questions_from_json(spec: Mapping[str, Any]) -> dict: | |
| """Wire-format questions ({key: {"type", "instructions", "criteria"}}, as the HTTP API takes; | |
| a top-level {"questions": {...}} wrapper is accepted) -> public Choice / Score / Noul.""" | |
| from jevlike.types import Choice, Noul, Score | |
| spec = spec.get("questions", spec) if isinstance(spec, Mapping) else spec | |
| out = {} | |
| for key, q in spec.items(): | |
| t = q.get("type") | |
| if t == "choice": | |
| out[key] = Choice(q["instructions"], q["criteria"]) | |
| elif t == "score": | |
| c = q["criteria"] | |
| levels = [f"{k}: {d}" if d else k for k, d in c.items()] if isinstance(c, dict) else list(c) | |
| out[key] = Score(q["instructions"], levels) | |
| elif t == "noul": | |
| out[key] = Noul(q["instructions"]) | |
| else: | |
| raise ValueError(f"question {key!r}: unknown type {t!r}") | |
| return out | |
| def main(argv=None) -> int: | |
| ap = argparse.ArgumentParser(description="Convert a document (PDF, image, HTML, text) to text and optionally " | |
| "answer SystemOne questions about it.") | |
| ap.add_argument("files", nargs="+", help="file(s); several files = pages of one document") | |
| ap.add_argument("--kind", choices=KINDS, default=None) | |
| ap.add_argument("--ocr", choices=OCR_ENGINES, default="tesseract") | |
| ap.add_argument("--lang", default="eng", help="tesseract language(s), e.g. eng+deu") | |
| ap.add_argument("--dpi", type=int, default=300) | |
| ap.add_argument("--min-chars", type=int, default=30) | |
| ap.add_argument("--no-tables", action="store_true") | |
| ap.add_argument("--no-preprocess", action="store_true") | |
| ap.add_argument("--keep-headers", action="store_true", help="do not strip repeated headers/footers") | |
| ap.add_argument("--max-pages", type=int, default=None) | |
| ap.add_argument("--json", action="store_true", help="print the IngestResult (and answers) as JSON") | |
| ap.add_argument("--questions", default=None, help="questions JSON (HTTP API format) to answer about the document") | |
| ap.add_argument("--checkpoint", default=None, help="checkpoint dir for --questions (e.g. checkpoints/jevlike-large-nf)") | |
| ap.add_argument("--device", default=None) | |
| args = ap.parse_args(argv) | |
| kw = dict(kind=args.kind, ocr=args.ocr, lang=args.lang, dpi=args.dpi, min_chars=args.min_chars, | |
| tables=not args.no_tables, preprocess=not args.no_preprocess, strip_headers=not args.keep_headers, | |
| max_pages=args.max_pages) | |
| src = args.files if len(args.files) > 1 else args.files[0] | |
| if args.questions: | |
| if not args.checkpoint: | |
| ap.error("--questions needs --checkpoint") | |
| from jevlike.predict import SystemOne | |
| qs = questions_from_json(json.loads(Path(args.questions).read_text())) | |
| model = SystemOne.load(args.checkpoint, device=args.device) | |
| res = predict_file(model, src, qs, ingest_kwargs=kw) | |
| print(json.dumps(res, indent=1, ensure_ascii=False, default=float)) | |
| return 0 | |
| r = to_text(src, **kw) | |
| if args.json: | |
| print(json.dumps(r.to_dict(), indent=1, ensure_ascii=False)) | |
| else: | |
| sys.stdout.write(r.text + "\n") | |
| info = ", ".join(f"{m}" for m in r.method_per_page) | |
| print(f"[{r.kind}: {len(r.pages)} page(s) [{info}], {len(r.text)} chars, {r.seconds:.2f}s]" | |
| + "".join(f"\n[warning] {w}" for w in r.warnings), file=sys.stderr) | |
| return 0 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |