geointel/scripts/export_operator_yolo_dataset.py

"""Export operator real-data samples to a YOLO detection dataset.

This is an operator/runtime helper. It converts the explicit orthophoto + GRB
reference sample manifest into local YOLO images/labels for model experiments.
It does not call GeoIntel APIs, does not train automatically and does not fetch
new provider data.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import sys
from pathlib import Path
from typing import Any, Iterable

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,
    create_training_release_manifest,
)


DEFAULT_MANIFEST_PATH = Path("/app/storage/operator-data/operator_samples_manifest.json")
DEFAULT_OUTPUT_DIR = Path("/app/storage/operator-data/yolo-building-dataset")
rasterio: Any = None
Transformer: Any = None
Image: Any = None


def file_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 parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Export operator real-data samples to a YOLO detection dataset.",
    )
    parser.add_argument(
        "--manifest-path",
        type=Path,
        default=Path(os.environ.get("OPERATOR_SAMPLE_MANIFEST_PATH", DEFAULT_MANIFEST_PATH)),
        help="operator_samples_manifest.json created by prepare_operator_real_data_samples.py.",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path(os.environ.get("OPERATOR_YOLO_DATASET_DIR", DEFAULT_OUTPUT_DIR)),
        help="Output directory for images, labels, dataset.yaml and summary JSON.",
    )
    parser.add_argument(
        "--val-samples",
        default=os.environ.get("OPERATOR_YOLO_VAL_SAMPLES", "turnhout"),
        help="Comma/space separated sample slugs assigned to validation. Defaults to turnhout.",
    )
    parser.add_argument(
        "--review-audit",
        type=Path,
        required=True,
        help="Passed corpus audit with accepted human-review evidence for this frozen source manifest.",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Remove and recreate output-dir before exporting.",
    )
    return parser.parse_args()


def ensure_dependencies() -> None:
    global Image, Transformer, rasterio
    try:
        import rasterio as rasterio_module
        from PIL import Image as image_module
        from pyproj import Transformer as transformer_class
    except Exception as exc:  # pragma: no cover - runtime environment only.
        raise SystemExit(
            "export_operator_yolo_dataset.py requires rasterio, pyproj and Pillow. "
            "Run it inside the GeoIntel all-in-one container or an equivalent GIS Python environment."
        ) from exc
    rasterio = rasterio_module
    Transformer = transformer_class
    Image = image_module


def split_slugs(raw: str) -> set[str]:
    return {value.strip().lower() for value in raw.replace(",", " ").split() if value.strip()}


def iter_geometry_coords(geometry: dict[str, Any]) -> Iterable[tuple[float, float]]:
    geometry_type = geometry.get("type")
    coordinates = geometry.get("coordinates")
    if geometry_type == "Polygon":
        for ring in coordinates or []:
            for point in ring:
                if len(point) >= 2:
                    yield float(point[0]), float(point[1])
    elif geometry_type == "MultiPolygon":
        for polygon in coordinates or []:
            for ring in polygon:
                for point in ring:
                    if len(point) >= 2:
                        yield float(point[0]), float(point[1])


def resolve_manifest_path(raw: str, manifest_path: Path) -> Path:
    path = Path(raw)
    if path.exists():
        return path
    if raw.startswith("/app/"):
        relative = Path(raw.removeprefix("/app/"))
        candidates = [
            Path.cwd() / relative,
            manifest_path.resolve().parent.parent.parent / relative,
        ]
        for candidate in candidates:
            if candidate.exists():
                return candidate
    return path


def image_array_from_raster(dataset: Any) -> Any:
    import numpy as np

    data = dataset.read()
    if data.shape[0] == 1:
        rgb = np.repeat(data[:1], 3, axis=0)
    else:
        rgb = data[:3]
    rgb = np.moveaxis(rgb, 0, -1)
    if rgb.dtype != np.uint8:
        rgb_min = float(np.nanmin(rgb))
        rgb_max = float(np.nanmax(rgb))
        if rgb_max > rgb_min:
            rgb = ((rgb - rgb_min) / (rgb_max - rgb_min) * 255.0).clip(0, 255).astype("uint8")
        else:
            rgb = np.zeros(rgb.shape, dtype="uint8")
    return rgb


def yolo_boxes_for_reference(reference_path: Path, dataset: Any) -> list[str]:
    reference = json.loads(reference_path.read_text(encoding="utf-8-sig"))
    features = reference.get("features") or []
    transformer = Transformer.from_crs("EPSG:4326", dataset.crs, always_xy=True)
    width = dataset.width
    height = dataset.height
    labels: list[str] = []

    for feature in features:
        properties = feature.get("properties") or {}
        if properties.get("source_name") != "grb":
            continue
        if properties.get("reference_layer_name") != "buildings":
            continue
        coords = list(iter_geometry_coords(feature.get("geometry") or {}))
        if not coords:
            continue
        xs, ys = zip(*(transformer.transform(lon, lat) for lon, lat in coords), strict=False)
        rows_cols = [dataset.index(x, y) for x, y in zip(xs, ys, strict=False)]
        rows = [row for row, _ in rows_cols]
        cols = [col for _, col in rows_cols]
        min_col = max(0, min(cols))
        max_col = min(width - 1, max(cols))
        min_row = max(0, min(rows))
        max_row = min(height - 1, max(rows))
        box_width = max_col - min_col
        box_height = max_row - min_row
        if box_width < 2 or box_height < 2:
            continue
        x_center = (min_col + max_col) / 2.0 / width
        y_center = (min_row + max_row) / 2.0 / height
        norm_width = box_width / width
        norm_height = box_height / height
        labels.append(f"0 {x_center:.8f} {y_center:.8f} {norm_width:.8f} {norm_height:.8f}")
    return labels


def export_sample(sample: dict[str, Any], manifest_path: Path, output_dir: Path, val_slugs: set[str]) -> dict[str, Any]:
    sample_slug = str(sample["sample_slug"])
    split = "val" if sample_slug.lower() in val_slugs else "train"
    raster_path = resolve_manifest_path(str(sample["raster_path"]), manifest_path)
    reference_path = resolve_manifest_path(str(sample["reference_path"]), manifest_path)
    if not raster_path.exists():
        raise SystemExit(f"Raster path is not readable for sample {sample_slug}: {raster_path}")
    if not reference_path.exists():
        raise SystemExit(f"Reference path is not readable for sample {sample_slug}: {reference_path}")

    image_path = output_dir / "images" / split / f"{sample_slug}.png"
    label_path = output_dir / "labels" / split / f"{sample_slug}.txt"
    image_path.parent.mkdir(parents=True, exist_ok=True)
    label_path.parent.mkdir(parents=True, exist_ok=True)

    with rasterio.open(raster_path) as dataset:
        image_array = image_array_from_raster(dataset)
        Image.fromarray(image_array).save(image_path)
        labels = yolo_boxes_for_reference(reference_path, dataset)
    label_path.write_text("\n".join(labels) + ("\n" if labels else ""), encoding="utf-8")
    return {
        "sample_slug": sample_slug,
        "split": split,
        "image_path": str(image_path),
        "label_path": str(label_path),
        "label_count": len(labels),
        "reference_feature_count": sample.get("reference_feature_count"),
        "raster_path": str(raster_path),
        "reference_path": str(reference_path),
    }


def write_dataset_yaml(output_dir: Path) -> Path:
    yaml_path = output_dir / "dataset.yaml"
    yaml_path.write_text(
        "\n".join(
            [
                f"path: {output_dir}",
                "train: images/train",
                "val: images/val",
                "names:",
                "  0: building",
                "",
            ]
        ),
        encoding="utf-8",
    )
    return yaml_path


def ensure_yolo_directories(output_dir: Path) -> None:
    for relative_path in ("images/train", "labels/train", "images/val", "labels/val"):
        (output_dir / relative_path).mkdir(parents=True, exist_ok=True)


def main() -> int:
    args = parse_args()
    manifest = json.loads(args.manifest_path.read_text(encoding="utf-8-sig"))
    try:
        assert_frozen_manifest_training_eligible(args.manifest_path, verify_live=True)
    except TrainingEligibilityError as exc:
        raise SystemExit(str(exc)) from exc
    ensure_dependencies()
    if args.force and args.output_dir.exists():
        shutil.rmtree(args.output_dir)
    args.output_dir.mkdir(parents=True, exist_ok=True)
    ensure_yolo_directories(args.output_dir)
    samples = manifest.get("samples") or []
    if not samples:
        raise SystemExit("Operator sample manifest contains no samples")
    val_slugs = split_slugs(args.val_samples)
    exported = [export_sample(sample, args.manifest_path, args.output_dir, val_slugs) for sample in samples]
    if not any(item["split"] == "train" for item in exported):
        raise SystemExit("YOLO dataset export produced no training samples")
    if not any(item["split"] == "val" for item in exported):
        raise SystemExit("YOLO dataset export produced no validation samples")
    dataset_yaml = write_dataset_yaml(args.output_dir)
    try:
        release_paths = create_training_release_manifest(
            train_yaml=dataset_yaml,
            corpus_manifest=args.manifest_path,
            review_audit_path=args.review_audit,
        )
    except TrainingReleaseError as exc:
        raise SystemExit(str(exc)) from exc
    summary = {
        "status": "ok",
        "dataset_yaml": str(dataset_yaml),
        "training_release_manifest": str(release_paths["release_manifest"]),
        "training_release_manifest_sha256": file_sha256(release_paths["release_manifest"]),
        "training_release_freeze": str(release_paths["release_freeze"]),
        "training_asset_manifest": str(release_paths["asset_manifest"]),
        "output_dir": str(args.output_dir),
        "class_names": ["building"],
        "sample_count": len(exported),
        "source_manifest": str(args.manifest_path),
        "source_manifest_sha256": file_sha256(args.manifest_path),
        "train_sample_count": sum(1 for item in exported if item["split"] == "train"),
        "val_sample_count": sum(1 for item in exported if item["split"] == "val"),
        "label_count": sum(item["label_count"] for item in exported),
        "samples": exported,
    }
    summary_path = args.output_dir / "yolo_dataset_summary.json"
    summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True), encoding="utf-8")
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    sys.exit(main())