geointel/scripts/compose_experimental_yolo_rehearsal_dataset.py

#!/usr/bin/env python3
"""Compose legacy rehearsal and fresh remediation tiles for an experiment."""

from __future__ import annotations

import argparse
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 export_experimental_yolo_train_shard import EXPERIMENTAL_ROOT, file_sha256  # noqa: E402


def _inside(path: Path, root: Path) -> bool:
    try:
        path.resolve().relative_to(root.resolve())
        return True
    except ValueError:
        return False


def validate_sources(rehearsal_summary: Path, remediation_summary: Path, output_dir: Path) -> tuple[dict, dict]:
    if not _inside(output_dir, EXPERIMENTAL_ROOT):
        raise ValueError(f"output must remain below {EXPERIMENTAL_ROOT}")
    rehearsal = json.loads(rehearsal_summary.read_text(encoding="utf-8"))
    remediation = json.loads(remediation_summary.read_text(encoding="utf-8"))
    marker_path = remediation_summary.parent / "EXPERIMENTAL_ONLY.json"
    if not marker_path.is_file():
        raise ValueError("fresh remediation source has no EXPERIMENTAL_ONLY marker")
    marker = json.loads(marker_path.read_text(encoding="utf-8"))
    if marker.get("promotion_allowed") is not False:
        raise ValueError("fresh remediation marker does not prohibit promotion")
    if remediation.get("source_manifest_sha256") != marker.get("source_manifest_sha256"):
        raise ValueError("fresh remediation marker/summary binding mismatch")
    rehearsal_tiles = [tile for tile in rehearsal.get("tiles") or [] if tile.get("kept", True)]
    remediation_tiles = [tile for tile in remediation.get("tiles") or [] if tile.get("kept", True)]
    if not rehearsal_tiles or not remediation_tiles:
        raise ValueError("both rehearsal and remediation summaries require retained tiles")
    if any(tile.get("split") not in {"train", "val"} for tile in rehearsal_tiles):
        raise ValueError("rehearsal source contains a protected or unsupported split")
    if any(tile.get("split") != "train" for tile in remediation_tiles):
        raise ValueError("fresh remediation source must be train-only")
    return rehearsal, remediation


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--rehearsal-summary", type=Path, required=True)
    parser.add_argument("--remediation-summary", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--remediation-repeat", type=int, default=2)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.remediation_repeat < 1 or args.remediation_repeat > 4:
        raise SystemExit("remediation-repeat must be between 1 and 4")
    if args.output_dir.exists():
        raise SystemExit(f"refusing to overwrite rehearsal dataset: {args.output_dir}")
    try:
        rehearsal, remediation = validate_sources(
            args.rehearsal_summary, args.remediation_summary, args.output_dir
        )
    except (OSError, ValueError, json.JSONDecodeError) as exc:
        raise SystemExit(str(exc)) from exc
    rehearsal_yaml = Path(str(rehearsal["dataset_yaml"]))
    remediation_yaml = Path(str(remediation["dataset_yaml"]))
    rehearsal_root = rehearsal_yaml.parent
    remediation_root = remediation_yaml.parent
    for directory in (
        rehearsal_root / "images" / "train",
        rehearsal_root / "images" / "val",
        remediation_root / "images" / "train",
    ):
        if not directory.is_dir():
            raise SystemExit(f"rehearsal dataset directory is missing: {directory}")
    args.output_dir.mkdir(parents=True)
    train_sources = [str(rehearsal_root / "images" / "train")] + [
        str(remediation_root / "images" / "train")
    ] * args.remediation_repeat
    yaml_path = args.output_dir / "dataset.yaml"
    yaml_path.write_text(
        "\n".join(
            [
                "path: /",
                "train:",
                *[f"  - {path}" for path in train_sources],
                f"val: {rehearsal_root / 'images' / 'val'}",
                "names:",
                "  0: building",
                "",
            ]
        ),
        encoding="utf-8",
    )
    summary = {
        "schema_version": 1,
        "status": "experimental_rehearsal_only",
        "promotion_allowed": False,
        "release_claim_allowed": False,
        "human_review_pending": True,
        "dataset_yaml": str(yaml_path),
        "dataset_yaml_sha256": file_sha256(yaml_path),
        "rehearsal_summary": str(args.rehearsal_summary),
        "rehearsal_summary_sha256": file_sha256(args.rehearsal_summary),
        "remediation_summary": str(args.remediation_summary),
        "remediation_summary_sha256": file_sha256(args.remediation_summary),
        "remediation_repeat": args.remediation_repeat,
        "rehearsal_train_tile_count": int(rehearsal.get("train_tile_count") or 0),
        "remediation_unique_train_tile_count": int(remediation.get("tile_count") or 0),
        "protected_v72_used": False,
    }
    (args.output_dir / "experimental_rehearsal_summary.json").write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    marker = {
        key: summary[key]
        for key in (
            "schema_version",
            "status",
            "promotion_allowed",
            "release_claim_allowed",
            "dataset_yaml_sha256",
            "rehearsal_summary_sha256",
            "remediation_summary_sha256",
        )
    }
    (args.output_dir / "EXPERIMENTAL_ONLY.json").write_text(
        json.dumps(marker, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(json.dumps(summary, indent=2))
    return 0


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