feat: harden governed PyTorch training programme

scripts/export_operator_yolo_tile_dataset.py

@@ -21,6 +21,9 @@ 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-tile-dataset")
+DEFAULT_CLASS_NAME = "building"
+DEFAULT_REFERENCE_SOURCE = "grb"
+DEFAULT_REFERENCE_LAYER = "buildings"
 REFERENCE_AOI_CATEGORY = "reference_aoi"
 PURE_EMPTY_BACKGROUND_CATEGORY = "pure_empty_negative"
 SPARSE_BACKGROUND_CATEGORY = "sparse_building_context"
@@ -74,6 +77,21 @@ def parse_args() -> argparse.Namespace:
         default=Path(os.environ.get("OPERATOR_YOLO_TILE_DATASET_DIR", DEFAULT_OUTPUT_DIR)),
         help="Output directory for tile images, labels, dataset.yaml and summary JSON.",
     )
+    parser.add_argument(
+        "--class-name",
+        default=os.environ.get("OPERATOR_YOLO_CLASS_NAME", DEFAULT_CLASS_NAME),
+        help="Canonical single detection class written to dataset.yaml.",
+    )
+    parser.add_argument(
+        "--reference-source",
+        default=os.environ.get("OPERATOR_YOLO_REFERENCE_SOURCE", DEFAULT_REFERENCE_SOURCE),
+        help="Required source_name in reference GeoJSON features.",
+    )
+    parser.add_argument(
+        "--reference-layer",
+        default=os.environ.get("OPERATOR_YOLO_REFERENCE_LAYER", DEFAULT_REFERENCE_LAYER),
+        help="Required reference_layer_name in reference GeoJSON features.",
+    )
     parser.add_argument("--tile-size", type=int, default=int(os.environ.get("OPERATOR_YOLO_TILE_SIZE", "256")))
     parser.add_argument("--stride", type=int, default=int(os.environ.get("OPERATOR_YOLO_TILE_STRIDE", "128")))
     parser.add_argument(
@@ -326,7 +344,14 @@ def iter_geometry_coords(geometry: dict[str, Any]) -> Iterable[tuple[float, floa
                         yield float(point[0]), float(point[1])
 
 
-def load_reference_pixel_boxes(reference_path: Path, dataset: Any, min_label_px: float) -> list[PixelBox]:
+def load_reference_pixel_boxes(
+    reference_path: Path,
+    dataset: Any,
+    min_label_px: float,
+    *,
+    reference_source: str,
+    reference_layer: str,
+) -> list[PixelBox]:
     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)
@@ -334,9 +359,9 @@ def load_reference_pixel_boxes(reference_path: Path, dataset: Any, min_label_px:
 
     for feature in features:
         properties = feature.get("properties") or {}
-        if properties.get("source_name") != "grb":
+        if str(properties.get("source_name") or "").strip().lower() != reference_source:
             continue
-        if properties.get("reference_layer_name") != "buildings":
+        if str(properties.get("reference_layer_name") or "").strip().lower() != reference_layer:
             continue
         coords = list(iter_geometry_coords(feature.get("geometry") or {}))
         if not coords:
@@ -423,7 +448,7 @@ def ensure_yolo_directories(output_dir: Path) -> None:
         (output_dir / relative_path).mkdir(parents=True, exist_ok=True)
 
 
-def write_dataset_yaml(output_dir: Path) -> Path:
+def write_dataset_yaml(output_dir: Path, class_name: str) -> Path:
     yaml_path = output_dir / "dataset.yaml"
     yaml_path.write_text(
         "\n".join(
@@ -432,7 +457,7 @@ def write_dataset_yaml(output_dir: Path) -> Path:
                 "train: images/train",
                 "val: images/val",
                 "names:",
-                "  0: building",
+                f"  0: {class_name}",
                 "",
             ]
         ),
@@ -454,6 +479,8 @@ def export_sample_tiles(
     background_negative_repeat: int,
     drop_low_variance_negatives: bool,
     blank_range_threshold: int,
+    reference_source: str,
+    reference_layer: str,
 ) -> list[dict[str, Any]]:
     sample_slug = str(sample["sample_slug"])
     sample_role = str(sample.get("sample_role") or "reference")
@@ -469,7 +496,13 @@ def export_sample_tiles(
 
     exported: list[dict[str, Any]] = []
     with rasterio.open(raster_path) as dataset:
-        boxes = load_reference_pixel_boxes(reference_path, dataset, min_label_px=min_label_px)
+        boxes = load_reference_pixel_boxes(
+            reference_path,
+            dataset,
+            min_label_px=min_label_px,
+            reference_source=reference_source,
+            reference_layer=reference_layer,
+        )
         for tile_index, tile_window in enumerate(iter_tile_windows(dataset.width, dataset.height, tile_size, stride)):
             labels = labels_for_tile(
                 tile_window,
@@ -566,6 +599,16 @@ def export_sample_tiles(
 
 def main() -> int:
     args = parse_args()
+    class_name = args.class_name.strip().lower()
+    reference_source = args.reference_source.strip().lower()
+    reference_layer = args.reference_layer.strip().lower()
+    for label, value in (
+        ("class-name", class_name),
+        ("reference-source", reference_source),
+        ("reference-layer", reference_layer),
+    ):
+        if not value or any(character not in "abcdefghijklmnopqrstuvwxyz0123456789_-" for character in value):
+            raise SystemExit(f"YOLO {label} must be a non-empty canonical slug")
     ensure_dependencies()
     if args.force and args.output_dir.exists():
         shutil.rmtree(args.output_dir)
@@ -597,6 +640,8 @@ def main() -> int:
                 background_negative_repeat=args.background_negative_repeat,
                 drop_low_variance_negatives=args.drop_low_variance_negatives,
                 blank_range_threshold=args.blank_range_threshold,
+                reference_source=reference_source,
+                reference_layer=reference_layer,
             )
         )
 
@@ -605,7 +650,7 @@ def main() -> int:
         raise SystemExit("YOLO tile dataset export produced no training tiles")
     if not any(tile["split"] == "val" for tile in kept_tiles):
         raise SystemExit("YOLO tile dataset export produced no validation tiles")
-    dataset_yaml = write_dataset_yaml(args.output_dir)
+    dataset_yaml = write_dataset_yaml(args.output_dir, class_name)
     positive_tiles = [tile for tile in kept_tiles if not tile["is_negative"]]
     negative_tiles = [tile for tile in kept_tiles if tile["is_negative"]]
     skipped_negative_tiles = [tile for tile in exported_tiles if not tile["kept"] and tile["is_negative"]]
@@ -619,7 +664,9 @@ def main() -> int:
         "status": "ok",
         "dataset_yaml": str(dataset_yaml),
         "output_dir": str(args.output_dir),
-        "class_names": ["building"],
+        "class_names": [class_name],
+        "reference_source": reference_source,
+        "reference_layer": reference_layer,
         "tile_size": args.tile_size,
         "stride": args.stride,
         "negative_keep_ratio": args.negative_keep_ratio,

scripts/train_operator_yolo_detector.sh

@@ -3,7 +3,7 @@ set -euo pipefail
 
 usage() {
   cat <<'EOF'
-Train a local operator YOLO building detector from an exported GeoIntel dataset.
+Train a local operator YOLO detector from an exported, governed GeoIntel dataset.
 
 Environment variables:
   OPERATOR_YOLO_DATASET_DIR   Directory containing dataset.yaml.
@@ -44,6 +44,7 @@ TRAIN_IMGSZ="${TRAIN_IMGSZ:-512}"
 TRAIN_BATCH="${TRAIN_BATCH:-2}"
 TRAIN_WORKERS="${TRAIN_WORKERS:-0}"
 TRAIN_DEVICE="${TRAIN_DEVICE:-cpu}"
+TRAIN_REQUIRE_CUDA="${TRAIN_REQUIRE_CUDA:-false}"
 if [[ -z "${PYTHON_BIN:-}" ]]; then
   if [[ -x "/opt/geointel/venv/bin/python" ]]; then
     PYTHON_BIN="/opt/geointel/venv/bin/python"
@@ -64,6 +65,7 @@ export TRAIN_IMGSZ
 export TRAIN_BATCH
 export TRAIN_WORKERS
 export TRAIN_DEVICE
+export TRAIN_REQUIRE_CUDA
 export PYTHON_BIN
 export SUMMARY_PATH
 
@@ -116,6 +118,7 @@ def seed_ultralytics_font() -> None:
 seed_ultralytics_font()
 
 from ultralytics import YOLO
+import torch
 
 dataset_yaml = Path(os.environ["DATASET_YAML"])
 base_model_path = Path(os.environ["YOLO_BASE_MODEL_PATH"])
@@ -128,6 +131,9 @@ image_size = int(os.environ["TRAIN_IMGSZ"])
 batch_size = int(os.environ["TRAIN_BATCH"])
 workers = int(os.environ["TRAIN_WORKERS"])
 device = os.environ["TRAIN_DEVICE"]
+require_cuda = os.environ["TRAIN_REQUIRE_CUDA"].strip().lower() in {"1", "true", "yes", "on"}
+if require_cuda and (not device.lower().startswith("cuda") or not torch.cuda.is_available()):
+    raise SystemExit("TRAIN_REQUIRE_CUDA is enabled but the requested CUDA device is unavailable")
 
 model = YOLO(str(base_model_path))
 model.train(
@@ -179,6 +185,9 @@ summary = {
     "batch_size": batch_size,
     "workers": workers,
     "device": device,
+    "cuda_required": require_cuda,
+    "cuda_available": torch.cuda.is_available(),
+    "torch_version": torch.__version__,
 }
 summary_path.parent.mkdir(parents=True, exist_ok=True)
 summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True), encoding="utf-8")