"""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 json import os import shutil import sys from pathlib import Path from typing import Any, Iterable 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 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( "--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() 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) manifest = json.loads(args.manifest_path.read_text(encoding="utf-8-sig")) 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) summary = { "status": "ok", "dataset_yaml": str(dataset_yaml), "output_dir": str(args.output_dir), "class_names": ["building"], "sample_count": len(exported), "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())