spectre_pdf 1.0.0

"""
Benchmark spectre_rs.Document.image_bytes() against pymupdf.Document.extract_image().

Speed (median of 3 warm calls) and accuracy (pixel-exact equality of the
decoded buffer in the image's native colorspace) on a fixed corpus of 23
PDFs spanning ICDAR 2013 and real IRS forms.

Byte-equality across libraries is unfair: pymupdf re-encodes raw-pixel
images with its own PNG defaults, and converts CMYK JPEGs to RGB. We
instead decode each side back to its native colorspace via Pillow and
compare pixel tensors, isolating fidelity from container choice. CMYK
images are reported separately since pymupdf converts to RGB (lossy
without an ICC profile) and spectre preserves the original CMYK bytes.

Usage:
    pip install pymupdf Pillow spectre_rs
    python scripts/bench_image_bytes_vs_pymupdf.py [--corpus DIR]
"""

from __future__ import annotations

import argparse
import glob
import io
import os
import statistics
import time
import urllib.request
import zipfile
from pathlib import Path

import fitz
import spectre_rs
from PIL import Image

ICDAR_URL = ("https://codeload.github.com/liminghao1630/"
             "ICDAR_2013_table_evaluate/zip/refs/heads/master")


def fetch_icdar(corpus_dir: Path) -> list[Path]:
    corpus_dir.mkdir(parents=True, exist_ok=True)
    zip_path = corpus_dir / "icdar2013.zip"
    if not zip_path.exists():
        print("fetch ICDAR 2013 corpus")
        urllib.request.urlretrieve(ICDAR_URL, zip_path)
    extracted = corpus_dir / "ICDAR_2013_table_evaluate-master"
    if not extracted.exists():
        with zipfile.ZipFile(zip_path) as z:
            z.extractall(corpus_dir)
    pdf_dir = (extracted /
               "icdar2013-competition-dataset-with-gt" / "pdf")
    return sorted(pdf_dir.glob("*.pdf"))[:20]


def fetch_irs(corpus_dir: Path) -> list[Path]:
    corpus_dir.mkdir(parents=True, exist_ok=True)
    urls = {
        "i1040gi.pdf": "https://www.irs.gov/pub/irs-pdf/i1040gi.pdf",
        "f1040.pdf": "https://www.irs.gov/pub/irs-pdf/f1040.pdf",
        "f1099int.pdf": "https://www.irs.gov/pub/irs-pdf/f1099int.pdf",
    }
    out = []
    for name, url in urls.items():
        p = corpus_dir / name
        if not p.exists():
            print(f"fetch {name}")
            urllib.request.urlretrieve(url, p)
        out.append(p)
    return out


def decode_native(buf: bytes, w: int, h: int) -> tuple[str, bytes] | None:
    """Decode an image to its native colorspace. Returns (mode, bytes)
    where mode is `RGB`, `CMYK`, `L`, or similar."""
    try:
        im = Image.open(io.BytesIO(buf))
        if im.size == (w, h):
            return (im.mode, im.tobytes())
    except Exception:
        return None
    return None


def bench(pdfs: list[Path]) -> None:
    print("=" * 110)
    print(f"{'PDF':<32} {'imgs':>5} {'extracted':>10} {'pmu ms':>8} "
          f"{'spec ms':>9} {'speedup':>9} {'pixel match':>13}")
    print("=" * 110)
    total_s = total_p = 0.0
    total_imgs = total_extracted = 0
    pixel_match = pixel_compared = 0
    cmyk_passthrough = 0  # CMYK preserved by spectre; pymupdf converts to RGB
    for path in pdfs:
        b = path.read_bytes()
        sd = spectre_rs.Document(b)
        s_imgs = sd.images()
        if not s_imgs:
            continue
        spec_runs = []
        s_extracted: list[tuple[int, bytes, str, int, int]] = []
        for _ in range(3):
            t0 = time.perf_counter()
            s_extracted = []
            for img in s_imgs:
                db = sd.image_bytes(img.xref)
                if db is not None:
                    s_extracted.append(
                        (img.xref, db.bytes, db.ext, img.width, img.height)
                    )
            spec_runs.append((time.perf_counter() - t0) * 1000)
        pmu_runs = []
        pmu_extracted: list[tuple[int, bytes, str, int, int]] = []
        for _ in range(3):
            d = fitz.open(stream=b)
            t0 = time.perf_counter()
            pmu_extracted = []
            for img in s_imgs:
                try:
                    info = d.extract_image(img.xref)
                    pmu_extracted.append((
                        img.xref, info["image"], info.get("ext", ""),
                        info["width"], info["height"],
                    ))
                except Exception:
                    pass
            pmu_runs.append((time.perf_counter() - t0) * 1000)
            d.close()
        spec_ms = statistics.median(spec_runs)
        pmu_ms = statistics.median(pmu_runs)
        total_s += spec_ms
        total_p += pmu_ms
        total_imgs += len(s_imgs)
        total_extracted += len(s_extracted)
        pmu_by_x = {x: (bts, w, h) for x, bts, _, w, h in pmu_extracted}
        spec_by_x = {x: (bts, w, h) for x, bts, _, w, h in s_extracted}
        local_match = local_count = local_cmyk = 0
        for x, (sb, sw, sh) in spec_by_x.items():
            if x in pmu_by_x:
                pb, pw, ph = pmu_by_x[x]
                sp = decode_native(sb, sw, sh)
                pp = decode_native(pb, pw, ph)
                if sp is None or pp is None:
                    continue
                s_mode, s_bytes = sp
                p_mode, p_bytes = pp
                if s_mode != p_mode or s_mode == "CMYK":
                    # Two distinct cases tracked together: (a) spectre
                    # preserved a non-RGB colorspace where pymupdf
                    # converted; (b) both libraries kept CMYK but use
                    # different channel conventions (Adobe inverted CMYK
                    # vs spec CMYK). In neither case is pixel-byte
                    # equality the right fidelity metric; count
                    # separately as "colorspace preserved."
                    local_cmyk += 1
                    continue
                local_count += 1
                if s_bytes == p_bytes:
                    local_match += 1
        pixel_match += local_match
        pixel_compared += local_count
        cmyk_passthrough += local_cmyk
        speedup = pmu_ms / spec_ms if spec_ms else float("inf")
        print(f"{path.name:<32} {len(s_imgs):>5} {len(s_extracted):>10} "
              f"{pmu_ms:>8.2f} {spec_ms:>9.2f} {speedup:>8.2f}x "
              f"{local_match}/{local_count:>11}")
    print("=" * 110)
    print(f"TOTAL  speedup {total_p / total_s:.2f}x  "
          f"({total_p:.1f}ms / {total_s:.1f}ms)")
    print(f"  extracted (viewable container): "
          f"{total_extracted}/{total_imgs} = {total_extracted / max(total_imgs, 1):.3f}")
    print(f"  pixel-exact equality vs pymupdf: "
          f"{pixel_match}/{pixel_compared} = {pixel_match / max(pixel_compared, 1):.3f}")
    if cmyk_passthrough:
        print(f"  CMYK preserved (pymupdf converted to RGB, not counted as miss): "
              f"{cmyk_passthrough} images")


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--corpus", default="bench-corpus")
    args = p.parse_args()
    corpus = Path(args.corpus)
    pdfs = fetch_icdar(corpus / "icdar2013") + fetch_irs(corpus / "irs")
    bench(pdfs)


if __name__ == "__main__":
    main()