scripts/prepare_dataset.py

# Converts the HF DocLayNet-base dataset into the image/label layout
# Ultralytics expects. Everything that can go wrong here goes wrong
# silently - bad labels don't raise, training just quietly learns the
# wrong thing. Three things beyond the obvious conversion:
#   1. collapses repeated per-line boxes (a 40-region page can show up as
#      ~1,100 duplicate annotations otherwise)
#   2. drops degenerate boxes and counts them (zero-area -> NaN loss later)
#   3. can render a sample with decoded labels drawn on, so I can eyeball
#      that the class index mapping is actually right
#
# Usage:
#   python scripts/prepare_dataset.py --out data/doclaynet
#   python scripts/prepare_dataset.py --out data/doclaynet --verify 12
#   python scripts/prepare_dataset.py --out data/doclaynet --limit 20   # smoke test
from __future__ import annotations

import argparse
import json
import random
from collections import Counter
from pathlib import Path

import sys

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from app.constants import CLASS_NAMES, ID_TO_CLASS  # noqa: E402

HF_DATASET = "pierreguillou/DocLayNet-base"

# HF split names don't match Ultralytics' directory convention
SPLITS = {"train": "train", "validation": "validation", "test": "test"}


def coco_to_yolo(
    bbox: list[float], img_w: int, img_h: int
) -> tuple[float, float, float, float]:
    """DocLayNet gives [x, y, w, h] top-left in pixels. Ultralytics wants
    [x_centre, y_centre, w, h] normalised 0-1. Easy to get the corner-to-
    centre shift wrong without noticing - tested for that reason."""
    x, y, w, h = bbox
    x_centre = (x + w / 2) / img_w
    y_centre = (y + h / 2) / img_h
    return (x_centre, y_centre, w / img_w, h / img_h)


def is_valid_bbox(bbox: list[float], img_w: int, img_h: int) -> bool:
    """Rejects zero-area/out-of-bounds boxes before they turn into NaN
    losses several epochs in. Counts rejections rather than just dropping
    silently."""
    x, y, w, h = bbox
    if w <= 0 or h <= 0:
        return False
    if x < 0 or y < 0:
        return False
    if x + w > img_w or y + h > img_h:
        return False
    return True


def dedupe_annotations(
    bboxes: list[list[float]], categories: list[int]
) -> list[tuple[tuple[float, ...], int]]:
    """DocLayNet-base repeats a block's box once per text line inside it -
    a paragraph with 6 lines shows up as 6 identical entries. Dedupe on
    (box, category) rather than box alone, since two classes occasionally
    share the same extent. First-seen order kept so output is reproducible."""
    if len(bboxes) != len(categories):
        raise ValueError(
            f"bboxes and categories must be the same length, "
            f"got {len(bboxes)} and {len(categories)}. "
            "If this fires, the dataset schema is not what this script assumes."
        )

    seen: set[tuple[tuple[float, ...], int]] = set()
    unique: list[tuple[tuple[float, ...], int]] = []

    for bbox, category in zip(bboxes, categories):
        key = (tuple(float(v) for v in bbox), int(category))
        if key not in seen:
            seen.add(key)
            unique.append(key)

    return unique


def _write_split(
    dataset, split_name: str, out_dir: Path, limit: int | None
) -> dict:
    """Writes one split to disk, returns the counts (these end up in the
    memo, especially the dedupe ratio - "removed 96% of raw annotations"
    needs a number next to it or it sounds like something broke)."""
    image_dir = out_dir / "images" / split_name
    label_dir = out_dir / "labels" / split_name
    image_dir.mkdir(parents=True, exist_ok=True)
    label_dir.mkdir(parents=True, exist_ok=True)

    rows = dataset[split_name]
    total = len(rows) if limit is None else min(limit, len(rows))

    stats = {
        "images": 0,
        "raw_annotations": 0,
        "after_dedupe": 0,
        "dropped_invalid": 0,
        "empty_pages": 0,
        "class_counts": Counter(),
        "doc_category_counts": Counter(),
    }

    # per-image manifest for evaluate.py's per-category mAP breakdown, and
    # for measuring train/test source-PDF leakage
    manifest: list[dict] = []

    for index in range(total):
        row = rows[index]
        image = row["image"]
        img_w, img_h = image.size

        stats["raw_annotations"] += len(row["bboxes_block"])

        annotations = dedupe_annotations(row["bboxes_block"], row["categories"])
        stats["after_dedupe"] += len(annotations)

        lines = []
        for bbox, category in annotations:
            if not is_valid_bbox(list(bbox), img_w, img_h):
                stats["dropped_invalid"] += 1
                continue
            xc, yc, w, h = coco_to_yolo(list(bbox), img_w, img_h)
            lines.append(f"{category} {xc:.6f} {yc:.6f} {w:.6f} {h:.6f}")
            stats["class_counts"][ID_TO_CLASS.get(category, f"UNKNOWN_{category}")] += 1

        # blank page is legit, but also what a conversion bug looks like
        if not lines:
            stats["empty_pages"] += 1

        stem = f"{split_name}_{index:06d}"
        image.convert("RGB").save(image_dir / f"{stem}.png")
        (label_dir / f"{stem}.txt").write_text("\n".join(lines), encoding="utf-8")

        doc_category = row.get("doc_category", "unknown")
        stats["doc_category_counts"][doc_category] += 1
        stats["images"] += 1

        manifest.append({
            "stem": stem,
            "doc_category": doc_category,
            "source_pdf": row.get("original_filename", "unknown"),
            "num_regions": len(lines),
        })

        if stats["images"] % 500 == 0:
            print(f"  {split_name}: {stats['images']}/{total}", flush=True)

    (out_dir / f"manifest_{split_name}.json").write_text(
        json.dumps(manifest, indent=2), encoding="utf-8"
    )

    stats["class_counts"] = dict(stats["class_counts"])
    stats["doc_category_counts"] = dict(stats["doc_category_counts"])
    return stats


def _render_label_check(dataset, out_dir: Path, sample_size: int) -> None:
    """Draws decoded class labels on a sample of pages so I can eyeball
    them. My class ordering (0-indexed alphabetical) is inferred, not
    confirmed - if it's off by one, Table becomes Section-header
    everywhere and training won't complain about it. No automated test
    catches that, only looking does."""
    from PIL import ImageDraw

    from scripts._render_utils import load_label_font

    label_font = load_label_font(28)

    check_dir = out_dir / "label_check"
    check_dir.mkdir(parents=True, exist_ok=True)

    rows = dataset["train"]
    random.seed(42)
    indices = random.sample(range(len(rows)), min(sample_size, len(rows)))

    palette = [
        "#e6194b", "#3cb44b", "#ffe119", "#4363d8", "#f58231", "#911eb4",
        "#46f0f0", "#f032e6", "#bcf60c", "#fabebe", "#008080",
    ]

    for index in indices:
        row = rows[index]
        image = row["image"].convert("RGB")
        draw = ImageDraw.Draw(image)

        for bbox, category in dedupe_annotations(row["bboxes_block"], row["categories"]):
            x, y, w, h = bbox
            colour = palette[category % len(palette)]
            label = ID_TO_CLASS.get(category, "?")
            draw.rectangle([x, y, x + w, y + h], outline=colour, width=4)

            # solid background behind the label - plain text is invisible
            # on dark scanned content otherwise
            text_box = draw.textbbox((x, y), label, font=label_font)
            draw.rectangle(
                [text_box[0] - 2, text_box[1] - 2, text_box[2] + 2, text_box[3] + 2],
                fill=colour,
            )
            draw.text((x, y), label, font=label_font, fill="white")

        image.save(check_dir / f"check_{index:06d}.png")

    print(f"\nWrote {len(indices)} annotated pages to {check_dir}")
    print("LOOK AT THESE before training. If the box labelled 'Table' is not")
    print("drawn around a table, the class index base is wrong and everything")
    print("downstream will be quietly meaningless.\n")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", default="data/doclaynet", help="output root")
    parser.add_argument(
        "--limit", type=int, default=None,
        help="cap images per split - for smoke tests only, not for a real run",
    )
    parser.add_argument(
        "--verify", type=int, default=0,
        help="render this many annotated pages for visual label checking",
    )
    args = parser.parse_args()

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    from datasets import ClassLabel, Features, Sequence, Value
    from datasets import load_dataset
    from datasets.features import Image as HFImage

    # The dataset's own loading script declares these as int64, but the
    # real coordinates are floats (hit one directly: 139.664355). Newer
    # pyarrow refuses that lossy cast where older versions silently
    # floored it. Can't edit someone else's script, so override the
    # schema on load instead - copied field-for-field, only these two
    # fixed to float64.
    doclaynet_features = Features({
        "id": Value("string"),
        "texts": Sequence(Value("string")),
        "bboxes_block": Sequence(Sequence(Value("float64"))),
        "bboxes_line": Sequence(Sequence(Value("float64"))),
        "categories": Sequence(ClassLabel(names=CLASS_NAMES)),
        "image": HFImage(),
        "page_hash": Value("string"),
        "original_filename": Value("string"),
        "page_no": Value("int32"),
        "num_pages": Value("int32"),
        "original_width": Value("int32"),
        "original_height": Value("int32"),
        "coco_width": Value("int32"),
        "coco_height": Value("int32"),
        "collection": Value("string"),
        "doc_category": Value("string"),
    })

    print(f"Loading {HF_DATASET} (3.8 GB on first run, cached after)...", flush=True)
    try:
        # trust_remote_code=True: this repo ships a loading script rather
        # than static parquet. Without this it prompts for confirmation,
        # which hangs forever under Kaggle's non-interactive commit runs.
        # Checked what the script does first - it's just the author's own
        # DocLayNet -> HF datasets conversion, nothing else.
        dataset = load_dataset(HF_DATASET, trust_remote_code=True, features=doclaynet_features)
    except RuntimeError as error:
        # datasets>=4.0.0 dropped loading-script support entirely -
        # requirements.txt pins <4.0.0 for this reason, but if an
        # environment already has a newer version cached, this is what
        # they'll hit.
        if "no longer supported" in str(error):
            raise RuntimeError(
                f"{error}\n\n"
                "This dataset repo uses the old Hugging Face 'loading script' "
                "format, which datasets>=4.0.0 removed support for entirely. "
                "Fix: pip install \"datasets<4.0.0\" (already pinned in "
                "requirements.txt - your environment likely has a newer "
                "version cached from something else)."
            ) from error
        raise

    if args.verify:
        _render_label_check(dataset, out_dir, args.verify)

    report = {"dataset": HF_DATASET, "splits": {}}
    for hf_split, dir_name in SPLITS.items():
        print(f"Converting split '{hf_split}'...", flush=True)
        report["splits"][dir_name] = _write_split(dataset, hf_split, out_dir, args.limit)

    _write_data_yaml(out_dir)

    report_path = out_dir / "prep_report.json"
    report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")

    print(f"\nWrote {report_path}")
    for split, stats in report["splits"].items():
        raw, kept = stats["raw_annotations"], stats["after_dedupe"]
        shrink = (1 - kept / raw) * 100 if raw else 0
        print(
            f"  {split:11s} {stats['images']:5d} images | "
            f"{raw:7d} raw -> {kept:6d} annotations ({shrink:.1f}% were repeats) | "
            f"{stats['dropped_invalid']} invalid dropped"
        )


def _write_data_yaml(out_dir: Path) -> None:
    """Generates the Ultralytics config instead of hand-writing it, so
    class names can't drift from app/constants.py."""
    names = "\n".join(f"  {idx}: {name}" for idx, name in enumerate(CLASS_NAMES))
    yaml = (
        "# Generated by scripts/prepare_dataset.py - do not edit by hand.\n"
        "# Class names come from app/constants.py so the two cannot drift.\n"
        f"path: {out_dir.resolve().as_posix()}\n"
        "train: images/train\n"
        "val: images/validation\n"
        "test: images/test\n"
        "\n"
        "names:\n"
        f"{names}\n"
    )
    (out_dir / "doclaynet.yaml").write_text(yaml, encoding="utf-8")


if __name__ == "__main__":
    main()