215 lines
7.1 KiB
Python
215 lines
7.1 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
|
|
SUPPORTED_MODEL_SUFFIXES = {".pt", ".onnx", ".engine"}
|
|
ENV_KEYS = ("GEOINTEL_INSTALL_AI", "YOLO_ENABLED", "YOLO_MODEL_PATH")
|
|
|
|
|
|
def _candidate_paths(models_dir: Path) -> list[Path]:
|
|
if not models_dir.exists():
|
|
return []
|
|
return sorted(
|
|
path.resolve()
|
|
for path in models_dir.rglob("*")
|
|
if path.is_file() and path.suffix.lower() in SUPPORTED_MODEL_SUFFIXES
|
|
)
|
|
|
|
|
|
def _container_path(host_model_path: Path, models_dir: Path, container_model_dir: str) -> str:
|
|
relative = host_model_path.resolve().relative_to(models_dir.resolve())
|
|
base = container_model_dir.rstrip("/")
|
|
return f"{base}/{relative.as_posix()}" if relative.as_posix() else base
|
|
|
|
|
|
def _read_env_lines(env_file: Path) -> list[str]:
|
|
if not env_file.exists():
|
|
return []
|
|
return env_file.read_text(encoding="utf-8").splitlines()
|
|
|
|
|
|
def _update_env_file(env_file: Path, updates: dict[str, str]) -> None:
|
|
existing_lines = _read_env_lines(env_file)
|
|
seen: set[str] = set()
|
|
next_lines: list[str] = []
|
|
|
|
for line in existing_lines:
|
|
stripped = line.strip()
|
|
if not stripped or stripped.startswith("#") or "=" not in line:
|
|
next_lines.append(line)
|
|
continue
|
|
|
|
key = line.split("=", 1)[0].strip()
|
|
if key in updates:
|
|
next_lines.append(f"{key}={updates[key]}")
|
|
seen.add(key)
|
|
else:
|
|
next_lines.append(line)
|
|
|
|
for key, value in updates.items():
|
|
if key not in seen:
|
|
next_lines.append(f"{key}={value}")
|
|
|
|
env_file.parent.mkdir(parents=True, exist_ok=True)
|
|
env_file.write_text("\n".join(next_lines).rstrip() + "\n", encoding="utf-8")
|
|
|
|
|
|
def _base_payload(args: argparse.Namespace, candidates: Iterable[Path]) -> dict[str, object]:
|
|
return {
|
|
"models_dir": str(Path(args.models_dir).resolve()),
|
|
"env_file": str(Path(args.env_file).resolve()),
|
|
"candidates": [str(path) for path in candidates],
|
|
"selected_host_model_path": None,
|
|
"selected_container_model_path": None,
|
|
"env_updates": {},
|
|
"apply": args.apply,
|
|
"applied": False,
|
|
"docker_restart_required": False,
|
|
"will_download_models": False,
|
|
"will_run_inference": False,
|
|
}
|
|
|
|
|
|
def _emit(payload: dict[str, object], *, as_json: bool) -> None:
|
|
if as_json:
|
|
print(json.dumps(payload, indent=2, sort_keys=True))
|
|
return
|
|
|
|
print(f"status: {payload['status']}")
|
|
print(f"message: {payload['message']}")
|
|
if payload.get("selected_host_model_path"):
|
|
print(f"host model: {payload['selected_host_model_path']}")
|
|
print(f"container model: {payload['selected_container_model_path']}")
|
|
if payload.get("env_updates"):
|
|
print("env updates:")
|
|
for key, value in payload["env_updates"].items(): # type: ignore[union-attr]
|
|
print(f" {key}={value}")
|
|
|
|
|
|
def configure(args: argparse.Namespace) -> tuple[int, dict[str, object]]:
|
|
models_dir = Path(args.models_dir).resolve()
|
|
env_file = Path(args.env_file).resolve()
|
|
candidates = _candidate_paths(models_dir)
|
|
payload = _base_payload(args, candidates)
|
|
|
|
if args.model_file:
|
|
selected = Path(args.model_file).resolve()
|
|
if not selected.exists() or not selected.is_file():
|
|
payload.update(
|
|
{
|
|
"status": "model_not_found",
|
|
"message": "The requested YOLO model file does not exist.",
|
|
}
|
|
)
|
|
return 2, payload
|
|
if selected.suffix.lower() not in SUPPORTED_MODEL_SUFFIXES:
|
|
payload.update(
|
|
{
|
|
"status": "unsupported_model_file",
|
|
"message": "The requested YOLO model file type is not supported.",
|
|
}
|
|
)
|
|
return 4, payload
|
|
try:
|
|
selected.relative_to(models_dir)
|
|
except ValueError:
|
|
payload.update(
|
|
{
|
|
"status": "model_outside_models_dir",
|
|
"message": "The requested model must be inside the mounted models directory.",
|
|
}
|
|
)
|
|
return 4, payload
|
|
elif not candidates:
|
|
payload.update(
|
|
{
|
|
"status": "no_model_found",
|
|
"message": "No local YOLO model file was found. Place a .pt, .onnx or .engine file in the models directory first.",
|
|
}
|
|
)
|
|
return 2, payload
|
|
elif len(candidates) > 1:
|
|
payload.update(
|
|
{
|
|
"status": "multiple_models_found",
|
|
"message": "Multiple model files were found. Re-run with --model-file to select one explicitly.",
|
|
}
|
|
)
|
|
return 3, payload
|
|
else:
|
|
selected = candidates[0]
|
|
|
|
selected_container_path = _container_path(selected, models_dir, args.container_model_dir)
|
|
updates = {
|
|
"GEOINTEL_INSTALL_AI": "true",
|
|
"YOLO_ENABLED": "true",
|
|
"YOLO_MODEL_PATH": selected_container_path,
|
|
}
|
|
payload.update(
|
|
{
|
|
"status": "ready_to_apply",
|
|
"message": "A local YOLO model was selected. Re-run with --apply to update the environment file.",
|
|
"selected_host_model_path": str(selected),
|
|
"selected_container_model_path": selected_container_path,
|
|
"env_updates": updates,
|
|
}
|
|
)
|
|
|
|
if args.apply:
|
|
_update_env_file(env_file, updates)
|
|
payload.update(
|
|
{
|
|
"status": "applied",
|
|
"message": "Environment file updated. Rebuild or restart the GeoIntel container to activate YOLO.",
|
|
"applied": True,
|
|
"docker_restart_required": True,
|
|
}
|
|
)
|
|
|
|
return 0, payload
|
|
|
|
|
|
def parse_args(argv: list[str]) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Configure GeoIntel to use an existing local YOLO model without downloading or running inference."
|
|
)
|
|
parser.add_argument(
|
|
"--models-dir",
|
|
default=os.environ.get("GEOINTEL_MODELS_PATH", "models"),
|
|
help="Host directory that is mounted into the container as the model directory.",
|
|
)
|
|
parser.add_argument(
|
|
"--container-model-dir",
|
|
default="/app/models",
|
|
help="Container path where --models-dir is mounted.",
|
|
)
|
|
parser.add_argument(
|
|
"--env-file",
|
|
default=".env",
|
|
help="Environment file to update when --apply is provided.",
|
|
)
|
|
parser.add_argument(
|
|
"--model-file",
|
|
help="Explicit model file to select when multiple local models exist.",
|
|
)
|
|
parser.add_argument("--apply", action="store_true", help="Write the required YOLO env values to --env-file.")
|
|
parser.add_argument("--json", action="store_true", help="Emit machine-readable JSON.")
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
args = parse_args(argv or sys.argv[1:])
|
|
exit_code, payload = configure(args)
|
|
_emit(payload, as_json=args.json)
|
|
return exit_code
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|