geointel/scripts/export_yolo_diagnostic_evaluation_tiles.py

#!/usr/bin/env python3
"""Tile an explicitly training-prohibited calibration corpus for diagnostics."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

try:
    from scripts.export_operator_yolo_tile_dataset import (
        ensure_dependencies,
        ensure_yolo_directories,
        export_sample_tiles,
        file_sha256,
        write_dataset_yaml,
    )
except ModuleNotFoundError:  # direct execution from /app/scripts
    from export_operator_yolo_tile_dataset import (
        ensure_dependencies,
        ensure_yolo_directories,
        export_sample_tiles,
        file_sha256,
        write_dataset_yaml,
    )


def validate_diagnostic_manifest(manifest_path: Path) -> dict[str, Any]:
    manifest = json.loads(manifest_path.read_text(encoding="utf-8-sig"))
    if manifest.get("purpose") != "non_protected_diagnostic_evaluation":
        raise ValueError("manifest is not a diagnostic-evaluation corpus")
    eligibility = manifest.get("training_eligibility")
    if not isinstance(eligibility, dict) or eligibility.get("status") != (
        "not_eligible_evaluation_only"
    ):
        raise ValueError("diagnostic manifest must be explicitly training-ineligible")
    samples = manifest.get("samples")
    if not isinstance(samples, list) or not samples:
        raise ValueError("diagnostic manifest contains no samples")
    slugs: set[str] = set()
    for sample in samples:
        if not isinstance(sample, dict):
            raise ValueError("diagnostic samples must be objects")
        slug = str(sample.get("sample_slug") or "").strip()
        if not slug or slug in slugs:
            raise ValueError(f"missing or duplicate diagnostic sample slug: {slug}")
        if sample.get("split") != "calibration":
            raise ValueError(f"diagnostic sample is not calibration-only: {slug}")
        if sample.get("sample_role") not in {"positive", "background_candidate"}:
            raise ValueError(f"invalid diagnostic sample role: {slug}")
        slugs.add(slug)
    marker_path = manifest_path.parent / "NO_TRAINING.json"
    if not marker_path.is_file():
        raise ValueError("diagnostic corpus has no NO_TRAINING.json marker")
    marker = json.loads(marker_path.read_text(encoding="utf-8"))
    if marker.get("training_allowed") is not False:
        raise ValueError("NO_TRAINING marker does not prohibit training")
    if marker.get("manifest_sha256") != file_sha256(manifest_path):
        raise ValueError("NO_TRAINING marker is not bound to the diagnostic manifest")
    return manifest


def is_canonical_evaluation_window(window: dict[str, Any], tile_size: int) -> bool:
    return bool(
        int(window["row_off"]) % tile_size == 0
        and int(window["col_off"]) % tile_size == 0
        and int(window["height"]) == tile_size
        and int(window["width"]) == tile_size
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--tile-size", type=int, default=512)
    parser.add_argument("--stride", type=int, default=512)
    parser.add_argument("--min-label-px", type=float, default=4.0)
    parser.add_argument("--min-label-visible-ratio", type=float, default=0.35)
    args = parser.parse_args()
    if args.output_dir.exists():
        raise SystemExit(f"refusing to overwrite evaluation tiles: {args.output_dir}")
    if args.stride != args.tile_size:
        raise SystemExit("diagnostic evaluation tiles must be non-overlapping")
    manifest_path = args.manifest.expanduser().resolve(strict=True)
    try:
        manifest = validate_diagnostic_manifest(manifest_path)
    except (OSError, ValueError, json.JSONDecodeError) as exc:
        raise SystemExit(str(exc)) from exc

    ensure_dependencies()
    ensure_yolo_directories(args.output_dir)
    exported: list[dict[str, Any]] = []
    for sample in manifest["samples"]:
        exported.extend(
            export_sample_tiles(
                sample=sample,
                manifest_path=manifest_path,
                output_dir=args.output_dir,
                val_slugs={str(sample["sample_slug"]).lower()},
                tile_size=args.tile_size,
                stride=args.stride,
                negative_keep_ratio=1.0,
                min_label_px=args.min_label_px,
                min_label_visible_ratio=args.min_label_visible_ratio,
                background_negative_repeat=1,
                drop_low_variance_negatives=False,
                blank_range_threshold=8.0,
                reference_source=str(sample["reference_source"]),
                reference_layer=str(sample["reference_layer"]),
            )
        )
    kept: list[dict[str, Any]] = []
    excluded_edge_cover: list[dict[str, Any]] = []
    for row in exported:
        window = row["window"]
        is_canonical_grid = bool(
            row["kept"]
            and is_canonical_evaluation_window(window, args.tile_size)
        )
        if is_canonical_grid:
            kept.append(row)
            continue
        if row["kept"]:
            excluded_edge_cover.append(
                {
                    "sample_slug": row["sample_slug"],
                    "tile_index": row["tile_index"],
                    "window": window,
                    "reason": "overlapping_edge_cover_tile",
                }
            )
            Path(row["image_path"]).unlink(missing_ok=True)
            Path(row["label_path"]).unlink(missing_ok=True)
    if not kept or any(row["split"] != "val" for row in kept):
        raise SystemExit("diagnostic exporter produced invalid split membership")
    if any((args.output_dir / "images" / "train").iterdir()) or any(
        (args.output_dir / "labels" / "train").iterdir()
    ):
        raise SystemExit("diagnostic exporter wrote training data")

    dataset_yaml = write_dataset_yaml(args.output_dir, "building")
    summary = {
        "schema_version": 1,
        "status": "ok_evaluation_only",
        "purpose": "non_protected_diagnostic_evaluation",
        "training_allowed": False,
        "release_claim_allowed": False,
        "dataset_yaml": str(dataset_yaml),
        "dataset_yaml_sha256": file_sha256(dataset_yaml),
        "source_manifest": str(manifest_path),
        "source_manifest_sha256": file_sha256(manifest_path),
        "tile_size": args.tile_size,
        "stride": args.stride,
        "tile_count": len(kept),
        "label_count": sum(int(row["label_count"]) for row in kept),
        "excluded_overlapping_edge_cover_tile_count": len(excluded_edge_cover),
        "excluded_overlapping_edge_cover_tiles": excluded_edge_cover,
        "selected_sample_slugs": sorted(
            {str(row["sample_slug"]) for row in kept}
        ),
        "tiles": kept,
    }
    summary_path = args.output_dir / "yolo_tile_dataset_summary.json"
    summary_path.write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    marker = {
        "schema_version": 1,
        "reason": "non_protected_diagnostic_evaluation_only",
        "training_allowed": False,
        "release_claim_allowed": False,
        "source_manifest_sha256": summary["source_manifest_sha256"],
        "summary_sha256": file_sha256(summary_path),
    }
    (args.output_dir / "NO_TRAINING.json").write_text(
        json.dumps(marker, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(json.dumps({key: value for key, value in summary.items() if key != "tiles"}))
    return 0


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