Add operator YOLO training dataset tooling

2026-07-07 05:19:16 +02:00
parent 306fcd1b24
commit 31a2aa6138
10 changed files with 607 additions and 0 deletions
@@ -0,0 +1,131 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+usage() {
+  cat <<'EOF'
+Train a local operator YOLO building detector from an exported GeoIntel dataset.
+
+Environment variables:
+  OPERATOR_YOLO_DATASET_DIR   Directory containing dataset.yaml.
+                               Default: /app/storage/operator-data/yolo-building-dataset
+  YOLO_BASE_MODEL_PATH         Existing local base .pt model path.
+                               Default: /app/models/yolov8n.pt
+  TRAIN_OUTPUT_DIR             Ultralytics project output directory.
+                               Default: /app/storage/training/operator-yolo
+  TRAIN_RUN_NAME               Ultralytics run name.
+                               Default: geointel-building-detector
+  TRAIN_MODEL_OUTPUT_PATH      Destination for the best trained .pt file.
+                               Default: /app/models/geointel-building-detector.pt
+  TRAIN_EPOCHS                 Training epochs. Default: 8
+  TRAIN_IMGSZ                  Image size. Default: 512
+  TRAIN_BATCH                  Batch size. Default: 2
+  TRAIN_WORKERS                Data-loader workers. Default: 0
+  TRAIN_DEVICE                 Device passed to Ultralytics. Default: cpu
+
+This helper is an operator/runtime smoke wrapper. It requires an existing local
+base model and an existing local dataset.yaml. It does not create app features.
+EOF
+}
+
+if [[ "${1:-}" == "--help" || "${1:-}" == "-h" ]]; then
+  usage
+  exit 0
+fi
+
+OPERATOR_YOLO_DATASET_DIR="${OPERATOR_YOLO_DATASET_DIR:-/app/storage/operator-data/yolo-building-dataset}"
+YOLO_BASE_MODEL_PATH="${YOLO_BASE_MODEL_PATH:-/app/models/yolov8n.pt}"
+TRAIN_OUTPUT_DIR="${TRAIN_OUTPUT_DIR:-/app/storage/training/operator-yolo}"
+TRAIN_RUN_NAME="${TRAIN_RUN_NAME:-geointel-building-detector}"
+TRAIN_MODEL_OUTPUT_PATH="${TRAIN_MODEL_OUTPUT_PATH:-/app/models/geointel-building-detector.pt}"
+TRAIN_EPOCHS="${TRAIN_EPOCHS:-8}"
+TRAIN_IMGSZ="${TRAIN_IMGSZ:-512}"
+TRAIN_BATCH="${TRAIN_BATCH:-2}"
+TRAIN_WORKERS="${TRAIN_WORKERS:-0}"
+TRAIN_DEVICE="${TRAIN_DEVICE:-cpu}"
+
+DATASET_YAML="${OPERATOR_YOLO_DATASET_DIR%/}/dataset.yaml"
+SUMMARY_PATH="${TRAIN_OUTPUT_DIR%/}/${TRAIN_RUN_NAME}/training_summary.json"
+export DATASET_YAML
+export YOLO_BASE_MODEL_PATH
+export TRAIN_OUTPUT_DIR
+export TRAIN_RUN_NAME
+export TRAIN_MODEL_OUTPUT_PATH
+export TRAIN_EPOCHS
+export TRAIN_IMGSZ
+export TRAIN_BATCH
+export TRAIN_WORKERS
+export TRAIN_DEVICE
+export SUMMARY_PATH
+
+if [[ ! -f "${DATASET_YAML}" ]]; then
+  echo "Dataset YAML not found: ${DATASET_YAML}" >&2
+  exit 1
+fi
+
+if [[ ! -f "${YOLO_BASE_MODEL_PATH}" ]]; then
+  echo "Base model file not found: ${YOLO_BASE_MODEL_PATH}" >&2
+  exit 1
+fi
+
+mkdir -p "${TRAIN_OUTPUT_DIR}" "$(dirname "${TRAIN_MODEL_OUTPUT_PATH}")"
+
+python - <<'PY'
+from __future__ import annotations
+
+import json
+import os
+import shutil
+from pathlib import Path
+
+from ultralytics import YOLO
+
+dataset_yaml = Path(os.environ["DATASET_YAML"])
+base_model_path = Path(os.environ["YOLO_BASE_MODEL_PATH"])
+train_output_dir = Path(os.environ["TRAIN_OUTPUT_DIR"])
+run_name = os.environ["TRAIN_RUN_NAME"]
+trained_model_output_path = Path(os.environ["TRAIN_MODEL_OUTPUT_PATH"])
+summary_path = Path(os.environ["SUMMARY_PATH"])
+epochs = int(os.environ["TRAIN_EPOCHS"])
+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"]
+
+model = YOLO(str(base_model_path))
+model.train(
+    data=str(dataset_yaml),
+    epochs=epochs,
+    imgsz=image_size,
+    batch=batch_size,
+    workers=workers,
+    device=device,
+    project=str(train_output_dir),
+    name=run_name,
+    exist_ok=True,
+    pretrained=True,
+    verbose=True,
+)
+
+best_path = train_output_dir / run_name / "weights" / "best.pt"
+if not best_path.exists():
+    raise SystemExit(f"Expected trained model artifact was not created: {best_path}")
+
+shutil.copy2(best_path, trained_model_output_path)
+summary = {
+    "status": "ok",
+    "dataset_yaml": str(dataset_yaml),
+    "base_model_path": str(base_model_path),
+    "train_output_dir": str(train_output_dir),
+    "train_run_name": run_name,
+    "trained_model_path": str(trained_model_output_path),
+    "best_artifact_path": str(best_path),
+    "epochs": epochs,
+    "image_size": image_size,
+    "batch_size": batch_size,
+    "workers": workers,
+    "device": device,
+}
+summary_path.parent.mkdir(parents=True, exist_ok=True)
+summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True), encoding="utf-8")
+print(json.dumps(summary, indent=2, sort_keys=True))
+PY