78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
import json
|
|
import os
|
|
import signal
|
|
import sys
|
|
import traceback
|
|
from collections.abc import Iterator, Sequence
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
|
|
# Force headless Matplotlib to avoid any thread/process GUI issues
|
|
os.environ["MPLBACKEND"] = "Agg"
|
|
|
|
from yolo_webui.config import TrainingConfig
|
|
from yolo_webui.trainer import TrainingEvent, TrainingRunner
|
|
|
|
|
|
@contextmanager
|
|
def _stop_signal_handlers(runner: TrainingRunner) -> Iterator[None]:
|
|
previous: dict[signal.Signals, signal.Handlers] = {}
|
|
|
|
def request_stop(_signum: int, _frame: object) -> None:
|
|
runner.request_stop()
|
|
|
|
for signum in (signal.SIGTERM, signal.SIGINT):
|
|
previous[signum] = signal.getsignal(signum)
|
|
signal.signal(signum, request_stop)
|
|
try:
|
|
yield
|
|
finally:
|
|
for signum, handler in previous.items():
|
|
signal.signal(signum, handler)
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
args = list(sys.argv[1:] if argv is None else argv)
|
|
if not args:
|
|
print(
|
|
"Usage: python -m yolo_webui.subprocess_runner <config_json_path>",
|
|
file=sys.stderr,
|
|
)
|
|
return 1
|
|
|
|
config_path = Path(args[0])
|
|
runner = TrainingRunner()
|
|
runner.prepare_run()
|
|
|
|
with _stop_signal_handlers(runner):
|
|
# The parent waits for this marker before sending a cooperative signal.
|
|
print("__YOLO_WEBUI_READY__", flush=True)
|
|
try:
|
|
with config_path.open("r", encoding="utf-8") as config_file:
|
|
config_dict = json.load(config_file)
|
|
config = TrainingConfig.from_dict(config_dict)
|
|
except Exception as exc:
|
|
print(f"Error loading config: {exc}", file=sys.stderr)
|
|
return 1
|
|
|
|
def handle_event(event: TrainingEvent) -> None:
|
|
event_dict = {
|
|
"kind": event.kind,
|
|
"message": event.message,
|
|
"epoch": event.epoch,
|
|
"total_epochs": event.total_epochs,
|
|
}
|
|
# Print structured JSON event so the parent process can parse it
|
|
print(f"__YOLO_WEBUI_EVENT__:{json.dumps(event_dict)}", flush=True)
|
|
|
|
try:
|
|
output_dir = runner.train(config, handle_event)
|
|
if output_dir:
|
|
print(f"__YOLO_WEBUI_RESULT__:{output_dir}", flush=True)
|
|
return 0
|
|
except Exception:
|
|
traceback.print_exc()
|
|
return 1
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|