Spaces:
Running
Running
| """ | |
| 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") | |