68 lines
2.4 KiB
Python
68 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
from typing import Any
|
|
|
|
from yolo_tui.config import MlflowConfig, TrainingConfig
|
|
from yolo_tui.trainer import TrainingEvent, TrainingRunner
|
|
|
|
|
|
def test_runner_wires_yolo_callbacks_and_returns_output(
|
|
monkeypatch: Any, tmp_path: Path
|
|
) -> None:
|
|
settings_updates: list[dict[str, bool]] = []
|
|
constructed: list[tuple[str, str]] = []
|
|
train_arguments: list[dict[str, Any]] = []
|
|
|
|
class FakeSettings:
|
|
def update(self, values: dict[str, bool]) -> None:
|
|
settings_updates.append(values)
|
|
|
|
class FakeYOLO:
|
|
def __init__(self, model: str, task: str) -> None:
|
|
constructed.append((model, task))
|
|
self.callbacks: dict[str, Any] = {}
|
|
self.trainer = SimpleNamespace(
|
|
args=SimpleNamespace(epochs=2),
|
|
epoch=0,
|
|
metrics={},
|
|
stop=False,
|
|
save_dir=tmp_path / "run",
|
|
)
|
|
|
|
def add_callback(self, name: str, callback: Any) -> None:
|
|
self.callbacks[name] = callback
|
|
|
|
def train(self, **kwargs: Any) -> None:
|
|
train_arguments.append(kwargs)
|
|
self.callbacks["on_train_start"](self.trainer)
|
|
for epoch in range(2):
|
|
self.trainer.epoch = epoch
|
|
self.trainer.metrics = {"metrics/mAP50": 0.5 + epoch / 10}
|
|
self.callbacks["on_train_epoch_end"](self.trainer)
|
|
self.callbacks["on_train_end"](self.trainer)
|
|
|
|
fake_ultralytics = ModuleType("ultralytics")
|
|
fake_ultralytics.YOLO = FakeYOLO # type: ignore[attr-defined]
|
|
fake_ultralytics.settings = FakeSettings() # type: ignore[attr-defined]
|
|
monkeypatch.setitem(sys.modules, "ultralytics", fake_ultralytics)
|
|
|
|
config = TrainingConfig(
|
|
dataset="dataset.yaml",
|
|
model="model.pt",
|
|
task="pose",
|
|
epochs=2,
|
|
mlflow=MlflowConfig(enabled=False),
|
|
)
|
|
events: list[TrainingEvent] = []
|
|
|
|
output = TrainingRunner().train(config, events.append)
|
|
|
|
assert output == tmp_path / "run"
|
|
assert constructed == [("model.pt", "pose")]
|
|
assert settings_updates == [{"mlflow": False}]
|
|
assert train_arguments[0]["data"] == "dataset.yaml"
|
|
assert train_arguments[0]["verbose"] is False
|
|
assert [event.kind for event in events] == ["info", "started", "epoch", "epoch", "success"]
|