256 lines
9.6 KiB
Python
256 lines
9.6 KiB
Python
"""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())
|