132 lines
4.4 KiB
Bash
132 lines
4.4 KiB
Bash
#!/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
|