from __future__ import annotations
import json
import statistics
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: str) -> list[str]:
return [t for t in text.lower().split() if t]
def doclaynet_partition(annotation_path: Path) -> tuple[list[str], list[int]]:
with open(annotation_path, "r", encoding="utf-8") as f:
data = json.load(f)
blocks: dict[int, list[str]] = defaultdict(list)
seen_categories: dict[int, str] = {}
for entry in data.get("form", []):
cat = entry.get("category", "")
if cat not in PARAGRAPH_CATEGORIES:
continue
block_id = entry.get("id_box")
text = entry.get("text", "")
if block_id is None or not text.strip():
continue
seen_categories.setdefault(block_id, cat)
blocks[block_id].append(text)
tokens: list[str] = []
block_ids: list[int] = []
for block_id, parts in blocks.items():
joined = " ".join(parts)
toks = tokenize(joined)
if not toks:
continue
tokens.extend(toks)
block_ids.extend([block_id] * len(toks))
return tokens, block_ids
def spectre_partition(pdf_bytes: bytes) -> tuple[list[str], list[int]]:
tokens: list[str] = []
block_ids: list[int] = []
doc = spectre_rs.Document(pdf_bytes)
for bid, block in enumerate(doc.blocks()):
toks = tokenize(block.text)
if not toks:
continue
tokens.extend(toks)
block_ids.extend([bid] * len(toks))
return tokens, block_ids
def pymupdf_partition(pdf_bytes: bytes) -> tuple[list[str], list[int]]:
tokens: list[str] = []
block_ids: list[int] = []
with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
for page in doc:
for bid, block in enumerate(page.get_text("blocks")):
if not isinstance(block[4], str):
continue
toks = tokenize(block[4])
if not toks:
continue
tokens.extend(toks)
block_ids.extend([bid] * len(toks))
return tokens, block_ids
def pair_f1(
cand_tokens: list[str],
cand_block_ids: list[int],
ref_tokens: list[str],
ref_block_ids: list[int],
) -> float:
ref_index: dict[str, list[int]] = defaultdict(list)
for pos, tok in enumerate(ref_tokens):
ref_index[tok].append(pos)
consumed: dict[str, int] = defaultdict(int)
aligned: list[tuple[int, int]] = []
for tok, cb in zip(cand_tokens, cand_block_ids):
positions = ref_index.get(tok)
if not positions:
continue
idx = consumed[tok]
if idx >= len(positions):
continue
ref_pos = positions[idx]
consumed[tok] += 1
aligned.append((cb, ref_block_ids[ref_pos]))
n = len(aligned)
if n < 2:
return 1.0 if n == 0 else 0.0
cand_groups: dict[int, int] = defaultdict(int)
ref_groups: dict[int, int] = defaultdict(int)
pair_counts: dict[tuple[int, int], int] = defaultdict(int)
for cb, rb in aligned:
cand_groups[cb] += 1
ref_groups[rb] += 1
pair_counts[(cb, rb)] += 1
def npairs(k: int) -> int:
return k * (k - 1) // 2
same_cand = sum(npairs(c) for c in cand_groups.values())
same_ref = sum(npairs(c) for c in ref_groups.values())
same_both = sum(npairs(c) for c in pair_counts.values())
tp = same_both
fp = same_cand - same_both
fn = same_ref - same_both
if tp == 0 and fp == 0 and fn == 0:
return 1.0
prec = tp / (tp + fp) if (tp + fp) else 0.0
rec = tp / (tp + fn) if (tp + fn) else 0.0
return (2 * prec * rec / (prec + rec)) if (prec + rec) else 0.0
def main() -> int:
if len(sys.argv) < 2:
print(__doc__, file=sys.stderr)
return 2
root = Path(sys.argv[1])
ann_dir = root / "annotations"
pdf_dir = root / "pdfs"
if not ann_dir.is_dir() or not pdf_dir.is_dir():
print(f"expected {root}/annotations and {root}/pdfs", file=sys.stderr)
return 2
annotations = sorted(ann_dir.glob("*.json"))
print(f"corpus: {len(annotations)} DocLayNet pages from {root}")
print(
f"{'page hash':<24}{'spectre F1':>14}{'pymupdf F1':>14}{'Δ':>8}"
)
print("-" * 60)
spectre_scores: list[float] = []
pymupdf_scores: list[float] = []
for i, ann in enumerate(annotations):
page_hash = ann.stem
pdf_path = pdf_dir / f"{page_hash}.pdf"
if not pdf_path.exists():
continue
try:
ref_t, ref_b = doclaynet_partition(ann)
except Exception as e: print(f"{page_hash[:20]:<24} ann err: {e}")
continue
if not ref_t:
continue
pdf_bytes = pdf_path.read_bytes()
try:
s_t, s_b = spectre_partition(pdf_bytes)
s_f1 = pair_f1(s_t, s_b, ref_t, ref_b)
except Exception as e: print(f"{page_hash[:20]:<24} spectre err: {e}")
continue
try:
p_t, p_b = pymupdf_partition(pdf_bytes)
p_f1 = pair_f1(p_t, p_b, ref_t, ref_b)
except Exception as e: print(f"{page_hash[:20]:<24} pymupdf err: {e}")
continue
delta = s_f1 - p_f1
spectre_scores.append(s_f1)
pymupdf_scores.append(p_f1)
if i < 30 or i % 10 == 0:
print(
f"{page_hash[:20]:<24}"
f"{s_f1:>14.3f}{p_f1:>14.3f}{delta:>+8.3f}"
)
print("-" * 60)
if spectre_scores:
print(
f"{'MEAN ({} pages)'.format(len(spectre_scores)):<24}"
f"{statistics.mean(spectre_scores):>14.3f}"
f"{statistics.mean(pymupdf_scores):>14.3f}"
f"{statistics.mean(spectre_scores) - statistics.mean(pymupdf_scores):>+8.3f}"
)
print(
f"{'MEDIAN':<24}"
f"{statistics.median(spectre_scores):>14.3f}"
f"{statistics.median(pymupdf_scores):>14.3f}"
f"{statistics.median(spectre_scores) - statistics.median(pymupdf_scores):>+8.3f}"
)
print()
print("Reference = DocLayNet human-annotated paragraph blocks (text /")
print("title / section-header / list-item / caption / footnote / page-")
print("header / page-footer). The first time we have an INDEPENDENT")
print("ground truth for spectre vs pymupdf — both libraries can be")
print("wrong simultaneously here, and the comparison is fair.")
return 0
if __name__ == "__main__":
sys.exit(main())