spectre_pdf 1.0.0

"""Failure-mode probe for DocLayNet worst-scoring pages.

Pulls human / pymupdf / spectre block partitions and dumps disagreement
examples for the pages where the libraries score lowest in
doclaynet_baseline.py.
"""
from __future__ import annotations

import json
import sys
from collections import defaultdict
from pathlib import Path

import fitz
import spectre_rs


PARAGRAPH_CATEGORIES = {
    "Text", "Title", "Section-header", "List-item",
    "Caption", "Footnote", "Page-header", "Page-footer",
}


def tokenize(text):
    return [t for t in text.lower().split() if t]


def human_blocks(ann_path):
    data = json.load(open(ann_path, "r", encoding="utf-8"))
    by_box = defaultdict(list)
    cat_by_box = {}
    for entry in data.get("form", []):
        bid = entry.get("id_box")
        if bid is None:
            continue
        txt = entry.get("text", "")
        cat = entry.get("category", "")
        if cat not in PARAGRAPH_CATEGORIES:
            continue
        by_box[bid].append(txt)
        cat_by_box.setdefault(bid, cat)
    out = []
    for bid in by_box:
        joined = " ".join(by_box[bid]).strip()
        if joined:
            out.append((bid, cat_by_box[bid], joined))
    return out


def spectre_blocks(pdf_bytes):
    doc = spectre_rs.Document(pdf_bytes)
    out = []
    for i, b in enumerate(doc.blocks()):
        if b.text.strip():
            out.append((i, b.text.strip()))
    return out


def pymupdf_blocks(pdf_bytes):
    out = []
    with fitz.open(stream=pdf_bytes, filetype="pdf") as d:
        bid = 0
        for p in d:
            for b in p.get_text("blocks"):
                if isinstance(b[4], str) and b[4].strip():
                    out.append((bid, b[4].strip()))
                    bid += 1
    return out


def find_containing(blocks, token, used):
    """Return index of first block whose text contains `token` (and not yet used)."""
    tok = token.lower()
    for i, b in enumerate(blocks):
        if i in used:
            continue
        text = b[-1].lower()
        if tok in text:
            return i
    return None


def show(label, block):
    text = block[-1].replace("\n", " ")
    if len(text) > 110:
        text = text[:107] + "..."
    print(f"    {label}: {text}")


def probe_page(root, page_hash):
    ann = next(root.glob(f"annotations/{page_hash}*.json"))
    pdf = next(root.glob(f"pdfs/{page_hash}*.pdf"))
    pdf_bytes = pdf.read_bytes()

    hum = human_blocks(ann)
    spc = spectre_blocks(pdf_bytes)
    pym = pymupdf_blocks(pdf_bytes)

    print()
    print(f"=== {page_hash} ===")
    cats = defaultdict(int)
    for _, c, _ in hum:
        cats[c] += 1
    print(f"  human: {len(hum)} blocks  cats={dict(cats)}")
    print(f"  pymupdf: {len(pym)} blocks")
    print(f"  spectre: {len(spc)} blocks")

    # Pick 3 example disagreements: first 3 human blocks with > 6 tokens,
    # find which pymupdf/spectre block contains the first distinctive token
    # of each. Distinctive = a longer (>5 char) token to avoid stopword
    # collisions.
    picked = 0
    used_pym = set()
    used_spc = set()
    for bid, cat, text in hum:
        toks = tokenize(text)
        big = [t for t in toks if len(t) > 5]
        if not big:
            continue
        anchor = big[0]
        pi = find_containing(pym, anchor, used_pym)
        si = find_containing(spc, anchor, used_spc)
        if pi is None or si is None:
            continue
        used_pym.add(pi); used_spc.add(si)
        print(f"  -- disagreement anchor={anchor!r} cat={cat}")
        show(f"HUM  [{bid}/{cat}]", (bid, cat, text))
        show(f"PYM  [{pi}]      ", pym[pi])
        show(f"SPC  [{si}]      ", spc[si])
        picked += 1
        if picked >= 3:
            break


def main():
    root = Path(sys.argv[1])
    for h in sys.argv[2:]:
        try:
            probe_page(root, h)
        except Exception as e:
            print(f"!! {h}: {e}")


if __name__ == "__main__":
    main()