geointel/scripts/build_regional_yolo_expert_dataset.py

#!/usr/bin/env python3
"""Create a leak-free train list for a regional YOLO expert."""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

from training_dataset_eligibility import (  # noqa: E402
    TrainingEligibilityError,
    assert_frozen_manifest_training_eligible,
)
from training_release_manifest import (  # noqa: E402
    TrainingReleaseError,
    assert_yolo_summary_bound_to_embedded_training_release,
    create_training_release_manifest,
)


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--summary", type=Path, required=True)
    parser.add_argument("--corpus-manifest", type=Path, required=True)
    parser.add_argument("--region", required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--positive-repeat", type=int, default=2)
    parser.add_argument("--negative-repeat", type=int, default=1)
    parser.add_argument(
        "--other-region-repeat",
        type=int,
        default=0,
        help="Include each train tile outside the expert region this many times.",
    )
    parser.add_argument(
        "--fixture-mode",
        action="store_true",
        help="Accept only an explicitly fixture-only corpus manifest; never use for operational training data.",
    )
    parser.add_argument(
        "--review-audit",
        type=Path,
        help="Accepted corpus-audit evidence; mandatory for an operational training release.",
    )
    args = parser.parse_args()
    if args.positive_repeat < 1 or args.negative_repeat < 1 or args.other_region_repeat < 0:
        raise SystemExit("regional repeats must be positive and other-region repeat non-negative")

    summary = json.loads(args.summary.read_text(encoding="utf-8"))
    manifest = json.loads(args.corpus_manifest.read_text(encoding="utf-8"))
    try:
        assert_frozen_manifest_training_eligible(
            args.corpus_manifest,
            fixture_mode=args.fixture_mode,
            verify_live=True,
        )
        assert_yolo_summary_bound_to_embedded_training_release(
            summary_path=args.summary,
            corpus_manifest=args.corpus_manifest,
            fixture_mode=args.fixture_mode,
        )
    except (TrainingEligibilityError, TrainingReleaseError) as exc:
        raise SystemExit(str(exc)) from exc
    samples = {item["sample_slug"]: item for item in manifest["samples"]}
    paths: list[str] = []
    selected_samples: set[str] = set()
    positive_tiles = negative_tiles = other_region_tiles = 0
    for tile in summary["tiles"]:
        sample = samples[tile["sample_slug"]]
        if sample["split"] != "train" or tile["split"] != "train":
            continue
        positive = int(tile.get("label_count") or 0) > 0
        if sample["region"] == args.region:
            repeat = args.positive_repeat if positive else args.negative_repeat
            positive_tiles += int(positive)
            negative_tiles += int(not positive)
        else:
            repeat = args.other_region_repeat
            other_region_tiles += int(repeat > 0)
        if repeat == 0:
            continue
        paths.extend([str(Path(tile["image_path"]).resolve())] * repeat)
        selected_samples.add(tile["sample_slug"])
    if not paths or not positive_tiles:
        raise SystemExit(f"no positive train tiles found for region {args.region!r}")

    args.output_dir.mkdir(parents=True, exist_ok=True)
    train_list = args.output_dir / "train.txt"
    train_list.write_text("\n".join(paths) + "\n", encoding="utf-8")
    dataset_yaml = args.output_dir / "dataset.yaml"
    dataset_yaml.write_text(
        f"path: {args.output_dir}\ntrain: {train_list}\n"
        f"val: {args.summary.parent / 'images' / 'val'}\nnames:\n  0: building\n",
        encoding="utf-8",
    )
    try:
        release_paths = create_training_release_manifest(
            train_yaml=dataset_yaml,
            corpus_manifest=args.corpus_manifest,
            review_audit_path=args.review_audit,
            fixture_mode=args.fixture_mode,
        )
    except TrainingReleaseError as exc:
        raise SystemExit(str(exc)) from exc
    evidence = {
        "schema_version": 1,
        "status": "ok",
        "region": args.region,
        "summary": str(args.summary),
        "summary_sha256": sha256(args.summary),
        "corpus_manifest": str(args.corpus_manifest),
        "corpus_manifest_sha256": sha256(args.corpus_manifest),
        "positive_repeat": args.positive_repeat,
        "negative_repeat": args.negative_repeat,
        "other_region_repeat": args.other_region_repeat,
        "source_positive_tile_count": positive_tiles,
        "source_negative_tile_count": negative_tiles,
        "source_other_region_tile_count": other_region_tiles,
        "sampled_train_entry_count": len(paths),
        "selected_train_samples": sorted(selected_samples),
        "protected_samples_in_training": [],
        "train_list": str(train_list),
        "dataset_yaml": str(dataset_yaml),
        "training_release_manifest": str(release_paths["release_manifest"]),
        "training_release_manifest_sha256": sha256(release_paths["release_manifest"]),
        "training_release_freeze": str(release_paths["release_freeze"]),
        "training_asset_manifest": str(release_paths["asset_manifest"]),
        "fixture_mode": bool(args.fixture_mode),
    }
    evidence_path = args.output_dir / "regional-expert-dataset.json"
    evidence_path.write_text(json.dumps(evidence, indent=2), encoding="utf-8")
    print(json.dumps(evidence, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())