Blame for core/layout_export.py

"""
core/layout_export.py — one annotation run, in whatever shape the next tool eats.

The annotator produces exact ground truth and writes it in its own format:
pages of elements, each with a label, a bbox in PDF points, and the text it
came from. That format is right for the tool that made it and wrong for every
tool that wants to consume it — COCO wants [x, y, w, h] and an integer category
table, YOLO wants normalised centres in per-image .txt files, VOC wants XML.

So this converts, and it converts carefully, because the failure mode here is
silent. A dataset with the wrong coordinate convention trains a model that is
confidently, uniformly wrong, and nothing in the pipeline complains. The two
places it can go wrong:

    origin      PDF measures `top` downward from the top edge, which happens to
                match image convention. Assuming otherwise flips every box
                vertically and looks plausible in aggregate statistics.
    units       points are not pixels. Anything rasterised at N DPI must scale
                by N/72, and a converter that ignores DPI produces boxes that
                are right only at 72.

Both are handled explicitly and asserted, rather than being left as something
the caller is expected to know.
"""
from __future__ import annotations

import json
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

SCHEMA = "signalmesh.layout_export/v1"

FORMATS = ("coco", "yolo", "voc", "hf", "jsonl")

PT_PER_INCH = 72.0

# `word` elements are the raw glyph runs; `region` elements are whole blocks.
# Training on both double-counts every headline, so a caller picks one.
GRAINS = ("region", "word", "both")


def _scale(dpi: float) -> float:
    """Points to pixels. At 72 DPI this is 1.0, which is exactly why a
    converter that forgets it appears to work."""
    return float(dpi) / PT_PER_INCH


def elements(page: Dict[str, Any], grain: str = "region") -> List[Dict[str, Any]]:
    els = page.get("elements") or []
    if grain == "both":
        return list(els)
    return [e for e in els if e.get("kind") == grain]


def labels_in(doc: Dict[str, Any], grain: str = "region") -> List[str]:
    """Every label actually present, not every label the palette defines.

    A category table full of classes with zero instances produces evaluation
    metrics that average in perfect scores for things the model never saw.
    """
    seen = {e.get("label") for p in doc.get("pages", [])
            for e in elements(p, grain) if e.get("label")}
    return sorted(seen)


def to_coco(doc: Dict[str, Any], *, grain: str = "region", dpi: float = 72.0,
            image_pattern: str = "{template}_p{page}.png") -> Dict[str, Any]:
    """COCO: xywh, one-based category ids, absolute pixels."""
    s = _scale(dpi)
    cats = labels_in(doc, grain)
    cat_id = {c: i + 1 for i, c in enumerate(cats)}
    images, anns = [], []
    ann_id = 1
    for page in doc.get("pages", []):
        n = page.get("page", 1)
        img_id = n
        images.append({
            "id": img_id,
            "file_name": image_pattern.format(template=doc.get("template", "doc"), page=n),
            "width": round(page.get("width", 0) * s),
            "height": round(page.get("height", 0) * s),
        })
        for e in elements(page, grain):
            x0, top, x1, bottom = [v * s for v in e["bbox"]]
            w, h = x1 - x0, bottom - top
            if w <= 0 or h <= 0:
                continue          # a zero-area box is not a training example
            anns.append({
                "id": ann_id, "image_id": img_id,
                "category_id": cat_id[e["label"]],
                "bbox": [round(x0, 2), round(top, 2), round(w, 2), round(h, 2)],
                "area": round(w * h, 2), "iscrowd": 0,
            })
            ann_id += 1
    return {
        "info": {"description": f"{doc.get('template','')} via docbank",
                 "profile": doc.get("profile"), "dpi": dpi, "grain": grain,
                 "fidelity_ok": bool((doc.get("fidelity") or {}).get("ok"))},
        "images": images, "annotations": anns,
        "categories": [{"id": i, "name": c} for c, i in cat_id.items()],
    }


def to_yolo(doc: Dict[str, Any], *, grain: str = "region") -> Dict[str, Any]:
    """YOLO: one text file per image, class + normalised cx cy w h.

    Normalised, so DPI is irrelevant here — which is the one place in this
    module where scaling would be a bug rather than a requirement.
    """
    cats = labels_in(doc, grain)
    idx = {c: i for i, c in enumerate(cats)}
    files: Dict[str, str] = {}
    for page in doc.get("pages", []):
        pw, ph = page.get("width", 0), page.get("height", 0)
        if not pw or not ph:
            continue
        lines = []
        for e in elements(page, grain):
            x0, top, x1, bottom = e["bbox"]
            w, h = (x1 - x0) / pw, (bottom - top) / ph
            cx, cy = (x0 + x1) / 2 / pw, (top + bottom) / 2 / ph
            if w <= 0 or h <= 0:
                continue
            lines.append(f"{idx[e['label']]} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}")
        name = f"{doc.get('template','doc')}_p{page.get('page',1)}.txt"
        files[name] = "\n".join(lines) + ("\n" if lines else "")
    return {"names": cats, "files": files,
            "data_yaml": "names:\n" + "".join(f"  {i}: {c}\n" for i, c in enumerate(cats))}


def to_voc(doc: Dict[str, Any], *, grain: str = "region", dpi: float = 72.0) -> Dict[str, str]:
    """Pascal VOC XML, one document per page."""
    from xml.sax.saxutils import escape
    s = _scale(dpi)
    out: Dict[str, str] = {}
    for page in doc.get("pages", []):
        n = page.get("page", 1)
        parts = [
            "<annotation>",
            f"  <folder>{escape(str(doc.get('template','doc')))}</folder>",
            f"  <filename>{escape(doc.get('template','doc'))}_p{n}.png</filename>",
            "  <size>",
            f"    <width>{round(page.get('width',0)*s)}</width>",
            f"    <height>{round(page.get('height',0)*s)}</height>",
            "    <depth>3</depth>", "  </size>",
        ]
        for e in elements(page, grain):
            x0, top, x1, bottom = [v * s for v in e["bbox"]]
            if x1 <= x0 or bottom <= top:
                continue
            parts += [
                "  <object>", f"    <name>{escape(str(e['label']))}</name>",
                "    <pose>Unspecified</pose><truncated>0</truncated><difficult>0</difficult>",
                "    <bndbox>",
                f"      <xmin>{round(x0)}</xmin><ymin>{round(top)}</ymin>",
                f"      <xmax>{round(x1)}</xmax><ymax>{round(bottom)}</ymax>",
                "    </bndbox>", "  </object>",
            ]
        parts.append("</annotation>")
        out[f"{doc.get('template','doc')}_p{n}.xml"] = "\n".join(parts)
    return out


def to_hf(doc: Dict[str, Any], *, grain: str = "region", dpi: float = 72.0) -> List[Dict[str, Any]]:
    """One row per page, objects in the column layout HF datasets expect."""
    s = _scale(dpi)
    cats = labels_in(doc, grain)
    idx = {c: i for i, c in enumerate(cats)}
    rows = []
    for page in doc.get("pages", []):
        objs = {"bbox": [], "category": [], "text": []}
        for e in elements(page, grain):
            x0, top, x1, bottom = [v * s for v in e["bbox"]]
            if x1 <= x0 or bottom <= top:
                continue
            objs["bbox"].append([round(x0, 2), round(top, 2),
                                 round(x1 - x0, 2), round(bottom - top, 2)])
            objs["category"].append(idx[e["label"]])
            objs["text"].append(e.get("text"))
        rows.append({
            "image_id": page.get("page", 1),
            "template": doc.get("template"), "profile": doc.get("profile"),
            "width": round(page.get("width", 0) * s),
            "height": round(page.get("height", 0) * s),
            "objects": objs,
            "categories": cats,
        })
    return rows


def to_jsonl(doc: Dict[str, Any], *, grain: str = "region", dpi: float = 72.0) -> str:
    return "\n".join(json.dumps(r, ensure_ascii=False)
                     for r in to_hf(doc, grain=grain, dpi=dpi))


def export(doc: Dict[str, Any], fmt: str, *, grain: str = "region",
           dpi: float = 72.0) -> Dict[str, Any]:
    """Convert, refusing anything that would produce a quietly-wrong dataset."""
    if fmt not in FORMATS:
        raise ValueError(f"unknown format {fmt!r} — have {', '.join(FORMATS)}")
    if grain not in GRAINS:
        raise ValueError(f"grain must be one of {GRAINS}, not {grain!r}")
    if not doc.get("pages"):
        raise ValueError("nothing to export — this run produced no pages")

    fid = doc.get("fidelity") or {}
    if fid and fid.get("ok") is False:
        raise ValueError(
            "this run failed the fidelity gate, so its boxes do not describe the "
            "page. Exporting it would put confidently-wrong labels into a "
            "dataset, which is worse than having no dataset.")

    body = {"coco": to_coco, "yolo": to_yolo, "voc": to_voc,
            "hf": to_hf, "jsonl": to_jsonl}[fmt](doc, grain=grain, **(
                {} if fmt == "yolo" else {"dpi": dpi}))
    counts: Dict[str, int] = {}
    for p in doc.get("pages", []):
        for e in elements(p, grain):
            counts[e["label"]] = counts.get(e["label"], 0) + 1
    return {"schema": SCHEMA, "format": fmt, "grain": grain, "dpi": dpi,
            "template": doc.get("template"), "labels": labels_in(doc, grain),
            "counts": counts, "pages": len(doc.get("pages", [])),
            "fidelity_ok": bool(fid.get("ok", True)), "data": body}


if __name__ == "__main__":
    # Shaped exactly like annotate.py writes it.
    doc = {
        "template": "lix_newspaper_class", "documentclass": "news",
        "profile": "news", "fidelity": {"ok": True, "max_drift_pt": 0.0},
        "labels": ["headline", "paragraph", "figure"],
        "pages": [{
            "page": 1, "width": 595.28, "height": 841.89,
            "elements": [
                {"label": "headline", "text": "Big News", "kind": "region",
                 "bbox": [72.0, 100.0, 523.0, 130.0]},
                {"label": "paragraph", "text": "body", "kind": "region",
                 "bbox": [72.0, 140.0, 300.0, 400.0]},
                {"label": "headline", "text": "Big", "kind": "word",
                 "bbox": [72.0, 100.0, 200.0, 130.0]},
                {"label": "figure", "text": None, "kind": "region",
                 "bbox": [10.0, 10.0, 10.0, 10.0]},        # zero area
            ],
            "equations": [],
        }],
    }

    # grain: regions and words are the same content twice
    assert len(elements(doc["pages"][0], "region")) == 3
    assert len(elements(doc["pages"][0], "word")) == 1
    assert len(elements(doc["pages"][0], "both")) == 4

    # only labels that actually occur — an empty class inflates evaluation
    assert labels_in(doc) == ["figure", "headline", "paragraph"]
    assert labels_in(doc, "word") == ["headline"]

    c = export(doc, "coco")["data"]
    assert len(c["images"]) == 1 and c["images"][0]["width"] == 595
    assert len(c["annotations"]) == 2, "the zero-area box must be dropped"
    a = c["annotations"][0]
    assert a["bbox"] == [72.0, 100.0, 451.0, 30.0], a["bbox"]   # xywh, not xyxy
    assert a["area"] == 13530.0
    assert {x["name"] for x in c["categories"]} == {"figure", "headline", "paragraph"}

    # points are not pixels
    c150 = export(doc, "coco", dpi=150)["data"]
    assert c150["images"][0]["width"] == round(595.28 * 150 / 72)
    assert c150["annotations"][0]["bbox"][0] == round(72.0 * 150 / 72, 2) == 150.0

    y = export(doc, "yolo")["data"]
    assert y["names"] == ["figure", "headline", "paragraph"]
    first = y["files"]["lix_newspaper_class_p1.txt"].splitlines()[0]
    cls, cx, cy, w, h = first.split()
    assert cls == "1"                                   # headline
    assert abs(float(cx) - (72 + 523) / 2 / 595.28) < 1e-6
    assert abs(float(w) - (523 - 72) / 595.28) < 1e-6
    assert all(0.0 <= float(v) <= 1.0 for v in (cx, cy, w, h)), "must be normalised"

    v = export(doc, "voc")["data"]["lix_newspaper_class_p1.xml"]
    assert "<name>headline</name>" in v and "<xmax>523</xmax>" in v
    assert v.count("<object>") == 2

    h = export(doc, "hf")["data"][0]
    assert h["objects"]["bbox"][0] == [72.0, 100.0, 451.0, 30.0]
    assert h["objects"]["category"][0] == 1 and h["objects"]["text"][0] == "Big News"

    j = export(doc, "jsonl")["data"]
    assert json.loads(j.splitlines()[0])["template"] == "lix_newspaper_class"

    meta = export(doc, "coco")
    assert meta["counts"] == {"headline": 1, "paragraph": 1, "figure": 1}
    assert meta["fidelity_ok"] is True and meta["pages"] == 1

    # ── the refusal that matters most ────────────────────────────────────────
    failed = dict(doc, fidelity={"ok": False, "error": "glyph count 812 vs 809"})
    try:
        export(failed, "coco")
        raise AssertionError("a failed fidelity gate must never reach a dataset")
    except ValueError as e:
        assert "worse than having no dataset" in str(e)

    for bad, why in (("parquet", "format"), ("coco", "grain")):
        try:
            export(doc, bad if why == "format" else "coco",
                   grain="region" if why == "format" else "glyphs")
            raise AssertionError(f"bad {why} must be refused")
        except ValueError:
            pass
    try:
        export({"pages": []}, "coco")
        raise AssertionError("an empty run must be refused")
    except ValueError as e:
        assert "no pages" in str(e)

    print("layout_export self-check OK")