187 lines
6.3 KiB
Bash
187 lines
6.3 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
|
|
PYTHON_BIN Python executable. Default: /opt/geointel/venv/bin/python
|
|
when present, otherwise python3.
|
|
|
|
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}"
|
|
if [[ -z "${PYTHON_BIN:-}" ]]; then
|
|
if [[ -x "/opt/geointel/venv/bin/python" ]]; then
|
|
PYTHON_BIN="/opt/geointel/venv/bin/python"
|
|
else
|
|
PYTHON_BIN="python3"
|
|
fi
|
|
fi
|
|
|
|
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 PYTHON_BIN
|
|
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_BIN}" - <<'PY'
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
|
|
def sha256_file(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def seed_ultralytics_font() -> None:
|
|
font_candidates = [
|
|
Path("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"),
|
|
Path("/usr/share/fonts/truetype/liberation2/LiberationSans-Regular.ttf"),
|
|
Path("/usr/share/fonts/truetype/freefont/FreeSans.ttf"),
|
|
]
|
|
source_font = next((font for font in font_candidates if font.exists()), None)
|
|
if source_font is None:
|
|
return
|
|
config_root = Path(os.environ.get("YOLO_CONFIG_DIR", str(Path.home() / ".config")))
|
|
target_font = config_root / "Ultralytics" / "Arial.ttf"
|
|
target_font.parent.mkdir(parents=True, exist_ok=True)
|
|
if not target_font.exists():
|
|
shutil.copy2(source_font, target_font)
|
|
|
|
|
|
seed_ultralytics_font()
|
|
|
|
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,
|
|
plots=False,
|
|
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)
|
|
dataset_summary_path = next(
|
|
(
|
|
candidate
|
|
for candidate in (
|
|
dataset_yaml.parent / "yolo_tile_dataset_summary.json",
|
|
dataset_yaml.parent / "yolo_dataset_summary.json",
|
|
)
|
|
if candidate.is_file()
|
|
),
|
|
None,
|
|
)
|
|
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),
|
|
"dataset_yaml_sha256": sha256_file(dataset_yaml),
|
|
"dataset_summary_path": str(dataset_summary_path) if dataset_summary_path else None,
|
|
"dataset_summary_sha256": sha256_file(dataset_summary_path) if dataset_summary_path else None,
|
|
"base_model_sha256": sha256_file(base_model_path),
|
|
"trained_model_sha256": sha256_file(trained_model_output_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
|