Compare commits
11 commits
weight-con
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
| c05cd91cde | |||
| 77a85bcb02 | |||
| 3640dff6d7 | |||
| ec017268bf | |||
| fed058d4c8 | |||
| d91eca366f | |||
| 212d03da67 | |||
| 1a8d1678d2 | |||
|
|
c87450393e | ||
|
|
53f758cc07 | ||
|
|
537ef2e489 |
30 changed files with 4199 additions and 497 deletions
17
.dockerignore
Normal file
17
.dockerignore
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
.git
|
||||
.agents
|
||||
.venv
|
||||
.pytest_cache
|
||||
**/__pycache__
|
||||
*.py[cod]
|
||||
.coverage
|
||||
build
|
||||
dist
|
||||
datasets
|
||||
models
|
||||
runs
|
||||
mlflow
|
||||
mlruns
|
||||
mlflow.db*
|
||||
*.pt
|
||||
.DS_Store
|
||||
20
Dockerfile
20
Dockerfile
|
|
@ -9,18 +9,22 @@ RUN pip install --no-cache-dir uv==0.10.6
|
|||
# Set working directory
|
||||
WORKDIR /workspace
|
||||
|
||||
# Install our dependencies into the system python without resolving/syncing
|
||||
# which would remove packages installed by ultralytics base image (like tensorrt).
|
||||
COPY pyproject.toml README.md ./
|
||||
RUN uv pip install --system fastapi mlflow uvicorn websockets
|
||||
ENV ULTRALYTICS_SAFE_LOAD=1
|
||||
|
||||
# Copy source code and install the project
|
||||
# Keep the training/GPU stack supplied by the Ultralytics image, while installing
|
||||
# the remaining application dependencies at the exact versions recorded in uv.lock.
|
||||
COPY pyproject.toml uv.lock README.md ./
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv export --locked --no-dev --group build --no-emit-project --prune ultralytics \
|
||||
--output-file /tmp/requirements.lock && \
|
||||
uv pip install --system --requirement /tmp/requirements.lock
|
||||
|
||||
# Copy source code and install only the project itself without re-resolving deps.
|
||||
COPY src ./src
|
||||
RUN uv pip install --system -e .
|
||||
RUN uv pip install --system --no-deps --no-build-isolation .
|
||||
|
||||
# Expose Web UI port and MLflow port
|
||||
# Expose Web UI port
|
||||
EXPOSE 8000
|
||||
EXPOSE 5000
|
||||
|
||||
# Start Web UI using the system entry point
|
||||
CMD ["yolo-train-webui", "--host", "0.0.0.0", "--port", "8000"]
|
||||
|
|
|
|||
48
README.md
48
README.md
|
|
@ -42,14 +42,17 @@ Ultralytics скачает веса. Пользовательские модел
|
|||
docker compose up --build
|
||||
```
|
||||
|
||||
WebUI будет доступен по `http://127.0.0.1:8000`. Compose намеренно публикует порт
|
||||
только на loopback. Не заменяйте адрес на `0.0.0.0` без аутентифицирующего reverse
|
||||
proxy: API позволяет запускать и останавливать ресурсоёмкие задачи.
|
||||
WebUI будет доступен по `http://127.0.0.1:8000`, а MLflow — по
|
||||
`http://127.0.0.1:5000`. Compose намеренно публикует оба порта только на loopback и
|
||||
хранит состояние MLflow в `./mlflow`. Не заменяйте адреса на `0.0.0.0` без
|
||||
аутентифицирующего reverse proxy: API позволяет запускать и останавливать
|
||||
ресурсоёмкие задачи.
|
||||
|
||||
Для NVIDIA GPU раскомментируйте секцию `deploy.resources.reservations.devices` в
|
||||
`docker-compose.yml`. Образ устанавливает зафиксированные в `uv.lock` зависимости;
|
||||
для другого варианта PyTorch используйте отдельно сгенерированный и проверенный
|
||||
lock-файл.
|
||||
`docker-compose.yml`. Web/MLflow-зависимости устанавливаются в версиях из
|
||||
`uv.lock`, а PyTorch, Ultralytics и GPU runtime предоставляет базовый образ
|
||||
Ultralytics. Для полностью воспроизводимой production-сборки дополнительно
|
||||
зафиксируйте базовый образ по digest.
|
||||
|
||||
## Разрешённые пути
|
||||
|
||||
|
|
@ -96,15 +99,46 @@ uv run mlflow ui --backend-store-uri sqlite:///mlflow.db
|
|||
артефакты запуска, а не версии MLflow Model Registry: raw YOLO checkpoint не имеет
|
||||
стандартной MLflow `MLmodel`-упаковки.
|
||||
|
||||
WebUI дополняет штатные метрики стабильным namespace `monitor/*`. Значения
|
||||
собираются после validation текущей эпохи, поэтому не запаздывают на одну эпоху:
|
||||
|
||||
| Задача | Основные дополнительные ряды |
|
||||
|---|---|
|
||||
| `detect` | box F1 при оптимальном confidence, mAP@0.75, метрики худшего класса |
|
||||
| `segment` | F1 при оптимальном confidence и mAP@0.75 отдельно для box и mask |
|
||||
| `classify` | top-1/top-5 error, macro precision/recall/F1, weighted F1, balanced accuracy |
|
||||
| `pose` | F1 при оптимальном confidence и mAP@0.75 отдельно для box и keypoints (OKS) |
|
||||
| `obb` | F1 при оптимальном confidence, mAP@0.75 и метрики худшего класса для oriented boxes |
|
||||
|
||||
Для всех задач также записываются исходный Ultralytics fitness и нормализованный
|
||||
`task_score` (для `segment`/`pose` сумма box+mask/keypoints приводится к диапазону
|
||||
0–1), суммарные train/validation loss, gap/ratio между ними, средний learning rate,
|
||||
время эпохи и скорость validation по стадиям. Теги `yolo.task` и
|
||||
`monitoring.schema_version` позволяют фильтровать совместимые запуски. Финальные
|
||||
per-class precision/recall/F1/AP и support сохраняются в артефакте
|
||||
`monitoring/task_metrics.csv`, чтобы не создавать сотни поэпоховых рядов для
|
||||
датасетов с большим числом классов. Monitor подключается внутри task-specific
|
||||
trainer и сохраняется при запуске Ultralytics DDP на нескольких GPU.
|
||||
|
||||
Шаги MLflow остаются совместимыми с Ultralytics: первая эпоха имеет `step=0`, а
|
||||
финальная проверка лучшего checkpoint добавляется отдельной точкой после последней
|
||||
эпохи. В WebUI эпохи отображаются привычно, начиная с 1.
|
||||
|
||||
Проверка интеграции на минимальных датасетах для всех пяти задач:
|
||||
|
||||
```bash
|
||||
uv run scripts/create_yolo26_smoke_datasets.py
|
||||
uv run scripts/run_yolo26_smoke_training.py --mlflow
|
||||
uv run scripts/verify_mlflow_smoke.py
|
||||
```
|
||||
|
||||
Smoke-runner присваивает всей пятёрке запусков уникальный тег `yolo.run_group`;
|
||||
проверка берёт только самый свежий batch и не смешивает его со старыми успешными
|
||||
run-ами.
|
||||
|
||||
Второй скрипт завершается с ошибкой, если отсутствует experiment/run, параметры,
|
||||
метрики, `results.csv`, `best.pt` или `last.pt` хотя бы для одной задачи.
|
||||
task-specific метрики, их история, теги, `monitoring/task_metrics.csv`,
|
||||
`results.csv`, `best.pt` или `last.pt` хотя бы для одной задачи.
|
||||
|
||||
## Проверка
|
||||
|
||||
|
|
|
|||
|
|
@ -3,21 +3,26 @@ services:
|
|||
build:
|
||||
context: .
|
||||
image: yolo-train-webui:latest
|
||||
network_mode: bridge
|
||||
ports:
|
||||
# The training API has no built-in user accounts, so expose it locally only.
|
||||
- "127.0.0.1:8000:8000"
|
||||
- "8000:8000"
|
||||
volumes:
|
||||
- ./datasets:/workspace/datasets
|
||||
- ./runs:/workspace/runs
|
||||
- ./models:/workspace/models
|
||||
- ./models/.config:/root/.config/Ultralytics
|
||||
environment:
|
||||
# Tracking URI for external MLflow server. Change host/IP if running MLflow on another machine/container.
|
||||
YOLO_WEBUI_MLFLOW_TRACKING_URI: http://host.docker.internal:5000
|
||||
extra_hosts:
|
||||
- "host.docker.internal:host-gateway"
|
||||
# Uncomment the block below on Linux with NVIDIA GPU to pass the graphics card into the container:
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: all
|
||||
capabilities: [gpu]
|
||||
# deploy:
|
||||
# resources:
|
||||
# reservations:
|
||||
# devices:
|
||||
# - driver: nvidia
|
||||
# count: all
|
||||
# capabilities: [gpu]
|
||||
ipc: host
|
||||
restart: unless-stopped
|
||||
|
|
|
|||
|
|
@ -16,13 +16,16 @@ dependencies = [
|
|||
yolo-train-webui = "yolo_webui.app:main"
|
||||
|
||||
[dependency-groups]
|
||||
build = [
|
||||
"hatchling==1.27.0",
|
||||
]
|
||||
dev = [
|
||||
"pytest>=8.3",
|
||||
"httpx",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
requires = ["hatchling==1.27.0"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
|
|
@ -31,4 +34,3 @@ packages = ["src/yolo_webui"]
|
|||
[tool.pytest.ini_options]
|
||||
addopts = "-q"
|
||||
testpaths = ["tests"]
|
||||
|
||||
|
|
|
|||
|
|
@ -4,8 +4,10 @@ from __future__ import annotations
|
|||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import traceback
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
from yolo_webui import TrainingConfig, TrainingRunner
|
||||
|
|
@ -39,6 +41,26 @@ def main() -> None:
|
|||
results: dict[str, dict[str, object]] = {}
|
||||
default_project = "runs/yolo26_mlflow_smoke/train" if args.mlflow else "runs/yolo26_smoke"
|
||||
project_dir = (args.project or Path(default_project)).resolve()
|
||||
smoke_batch_id = uuid.uuid4().hex if args.mlflow else ""
|
||||
if args.mlflow:
|
||||
import mlflow
|
||||
|
||||
project_dir.parent.mkdir(parents=True, exist_ok=True)
|
||||
if args.tracking_uri.startswith("sqlite:///"):
|
||||
tracking_db = Path(args.tracking_uri.removeprefix("sqlite:///"))
|
||||
tracking_db.parent.mkdir(parents=True, exist_ok=True)
|
||||
os.environ["YOLO_WEBUI_MLFLOW_RUN_GROUP"] = smoke_batch_id
|
||||
mlflow.set_tracking_uri(args.tracking_uri)
|
||||
mlflow.set_experiment(args.experiment)
|
||||
with mlflow.start_run(run_name=f"smoke-batch-{smoke_batch_id}") as anchor:
|
||||
mlflow.set_tags(
|
||||
{
|
||||
"smoke.anchor": "true",
|
||||
"yolo.run_group": smoke_batch_id,
|
||||
}
|
||||
)
|
||||
print(f"MLflow smoke batch: {smoke_batch_id}", flush=True)
|
||||
|
||||
for task, (dataset, model) in TASKS.items():
|
||||
print(f"\n=== {task}: {model} ===", flush=True)
|
||||
config = TrainingConfig(
|
||||
|
|
@ -83,14 +105,24 @@ def main() -> None:
|
|||
"error": f"{type(exc).__name__}: {exc}",
|
||||
}
|
||||
|
||||
summary_path = project_dir.parent / "smoke_summary.json" if args.mlflow else project_dir / "smoke_summary.json"
|
||||
summary_path = (
|
||||
project_dir.parent / "smoke_summary.json"
|
||||
if args.mlflow
|
||||
else project_dir / "smoke_summary.json"
|
||||
)
|
||||
summary_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
summary = {
|
||||
"smoke_batch_id": smoke_batch_id or None,
|
||||
"tracking_uri": args.tracking_uri if args.mlflow else None,
|
||||
"experiment": args.experiment if args.mlflow else None,
|
||||
"tasks": results,
|
||||
}
|
||||
summary_path.write_text(
|
||||
json.dumps(results, indent=2, ensure_ascii=False) + "\n",
|
||||
json.dumps(summary, indent=2, ensure_ascii=False) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
print(f"\nSummary: {summary_path.resolve()}")
|
||||
print(json.dumps(results, indent=2, ensure_ascii=False))
|
||||
print(json.dumps(summary, indent=2, ensure_ascii=False))
|
||||
if any(result["status"] != "succeeded" for result in results.values()):
|
||||
raise SystemExit(1)
|
||||
|
||||
|
|
|
|||
|
|
@ -3,15 +3,139 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import math
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import mlflow
|
||||
from mlflow.entities import Run
|
||||
from mlflow.tracking import MlflowClient
|
||||
from yolo_webui.mlflow_metrics import PER_CLASS_FIELDS
|
||||
|
||||
|
||||
TASKS = ("detect", "segment", "classify", "pose", "obb")
|
||||
REQUIRED_ARTIFACTS = {"weights/best.pt", "weights/last.pt", "results.csv"}
|
||||
REQUIRED_ARTIFACTS = {
|
||||
"weights/best.pt",
|
||||
"weights/last.pt",
|
||||
"results.csv",
|
||||
"monitoring/task_metrics.csv",
|
||||
}
|
||||
BOX_METRICS = {
|
||||
"metrics/precisionB",
|
||||
"metrics/recallB",
|
||||
"metrics/mAP50B",
|
||||
"metrics/mAP50-95B",
|
||||
}
|
||||
REQUIRED_METRICS_BY_TASK = {
|
||||
"detect": {
|
||||
*BOX_METRICS,
|
||||
"train/box_loss",
|
||||
"train/cls_loss",
|
||||
"train/dfl_loss",
|
||||
"val/box_loss",
|
||||
"val/cls_loss",
|
||||
"val/dfl_loss",
|
||||
"monitor/quality/box/f1_at_optimal_confidence",
|
||||
"monitor/quality/box/map75",
|
||||
},
|
||||
"segment": {
|
||||
*BOX_METRICS,
|
||||
"metrics/precisionM",
|
||||
"metrics/recallM",
|
||||
"metrics/mAP50M",
|
||||
"metrics/mAP50-95M",
|
||||
"train/seg_loss",
|
||||
"train/box_loss",
|
||||
"train/cls_loss",
|
||||
"train/dfl_loss",
|
||||
"val/seg_loss",
|
||||
"val/box_loss",
|
||||
"val/cls_loss",
|
||||
"val/dfl_loss",
|
||||
"monitor/quality/box/f1_at_optimal_confidence",
|
||||
"monitor/quality/mask/f1_at_optimal_confidence",
|
||||
"monitor/quality/mask/map75",
|
||||
},
|
||||
"classify": {
|
||||
"metrics/accuracy_top1",
|
||||
"metrics/accuracy_top5",
|
||||
"train/loss",
|
||||
"val/loss",
|
||||
"monitor/quality/classification/top1_error",
|
||||
"monitor/quality/classification/macro_f1",
|
||||
"monitor/quality/classification/balanced_accuracy",
|
||||
},
|
||||
"pose": {
|
||||
*BOX_METRICS,
|
||||
"metrics/precisionP",
|
||||
"metrics/recallP",
|
||||
"metrics/mAP50P",
|
||||
"metrics/mAP50-95P",
|
||||
"train/pose_loss",
|
||||
"train/box_loss",
|
||||
"train/kobj_loss",
|
||||
"train/cls_loss",
|
||||
"train/dfl_loss",
|
||||
"val/pose_loss",
|
||||
"val/box_loss",
|
||||
"val/kobj_loss",
|
||||
"val/cls_loss",
|
||||
"val/dfl_loss",
|
||||
"monitor/quality/box/f1_at_optimal_confidence",
|
||||
"monitor/quality/keypoints/f1_at_optimal_confidence",
|
||||
"monitor/quality/keypoints/map75",
|
||||
},
|
||||
"obb": {
|
||||
*BOX_METRICS,
|
||||
"train/box_loss",
|
||||
"train/cls_loss",
|
||||
"train/dfl_loss",
|
||||
"train/angle_loss",
|
||||
"val/box_loss",
|
||||
"val/cls_loss",
|
||||
"val/dfl_loss",
|
||||
"val/angle_loss",
|
||||
"monitor/quality/oriented_box/f1_at_optimal_confidence",
|
||||
"monitor/quality/oriented_box/map75",
|
||||
},
|
||||
}
|
||||
COMMON_MONITOR_METRICS = {
|
||||
"monitor/fitness/ultralytics_current",
|
||||
"monitor/fitness/ultralytics_best",
|
||||
"monitor/fitness/task_score_current",
|
||||
"monitor/fitness/task_score_best",
|
||||
"monitor/loss/train_total",
|
||||
"monitor/loss/validation_total",
|
||||
"monitor/loss/generalization_gap",
|
||||
"monitor/loss/validation_to_train_ratio",
|
||||
"monitor/optimization/learning_rate_mean",
|
||||
"monitor/performance/epoch_seconds",
|
||||
"monitor/performance/validation_preprocess_ms_per_image",
|
||||
"monitor/performance/validation_inference_ms_per_image",
|
||||
"monitor/performance/validation_loss_ms_per_image",
|
||||
"monitor/performance/validation_postprocess_ms_per_image",
|
||||
}
|
||||
HISTORY_METRIC_BY_TASK = {
|
||||
"detect": "monitor/quality/box/f1_at_optimal_confidence",
|
||||
"segment": "monitor/quality/mask/f1_at_optimal_confidence",
|
||||
"classify": "monitor/quality/classification/macro_f1",
|
||||
"pose": "monitor/quality/keypoints/f1_at_optimal_confidence",
|
||||
"obb": "monitor/quality/oriented_box/f1_at_optimal_confidence",
|
||||
}
|
||||
EXPECTED_COMPONENTS = {
|
||||
"detect": {"box"},
|
||||
"segment": {"box", "mask"},
|
||||
"classify": {"classification"},
|
||||
"pose": {"box", "keypoints"},
|
||||
"obb": {"oriented_box"},
|
||||
}
|
||||
|
||||
|
||||
def require(condition: object, message: str) -> None:
|
||||
if not condition:
|
||||
raise AssertionError(message)
|
||||
|
||||
|
||||
def artifact_paths(client: MlflowClient, run_id: str, path: str = "") -> set[str]:
|
||||
|
|
@ -24,12 +148,74 @@ def artifact_paths(client: MlflowClient, run_id: str, path: str = "") -> set[str
|
|||
return result
|
||||
|
||||
|
||||
def latest_task_run(runs: list[Run], task: str) -> Run:
|
||||
def latest_task_run(runs: list[Run], task: str, run_group: str) -> Run:
|
||||
expected_name = f"{task}-smoke"
|
||||
for run in runs:
|
||||
if run.data.tags.get("mlflow.runName") == expected_name:
|
||||
if (
|
||||
run.data.tags.get("mlflow.runName") == expected_name
|
||||
and run.data.tags.get("yolo.run_group") == run_group
|
||||
):
|
||||
return run
|
||||
raise AssertionError(f"MLflow run not found: {expected_name}")
|
||||
raise AssertionError(
|
||||
f"MLflow run not found: {expected_name}, yolo.run_group={run_group}"
|
||||
)
|
||||
|
||||
|
||||
def latest_smoke_run_group(runs: list[Run]) -> str:
|
||||
for run in runs:
|
||||
if (
|
||||
run.data.tags.get("smoke.anchor") == "true"
|
||||
and run.data.tags.get("yolo.run_group")
|
||||
):
|
||||
return run.data.tags["yolo.run_group"]
|
||||
raise AssertionError("No MLflow smoke batch with yolo.run_group tag found")
|
||||
|
||||
|
||||
def verify_per_class_artifact(
|
||||
client: MlflowClient,
|
||||
run_id: str,
|
||||
task: str,
|
||||
) -> int:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
downloaded = Path(
|
||||
client.download_artifacts(
|
||||
run_id,
|
||||
"monitoring/task_metrics.csv",
|
||||
dst_path=directory,
|
||||
)
|
||||
)
|
||||
with downloaded.open(encoding="utf-8") as stream:
|
||||
reader = csv.DictReader(stream)
|
||||
rows = list(reader)
|
||||
require(
|
||||
tuple(reader.fieldnames or ()) == PER_CLASS_FIELDS,
|
||||
f"Invalid task_metrics.csv header for {task}: {reader.fieldnames}",
|
||||
)
|
||||
|
||||
require(rows, f"Empty task_metrics.csv for {task}")
|
||||
components = {row["component"] for row in rows}
|
||||
require(
|
||||
components == EXPECTED_COMPONENTS[task],
|
||||
f"Invalid task_metrics.csv components for {task}: {sorted(components)}",
|
||||
)
|
||||
for row in rows:
|
||||
require(row["task"] == task, f"Invalid task in task_metrics.csv: {row}")
|
||||
require(row["class_id"].isdigit(), f"Invalid class_id in task_metrics.csv: {row}")
|
||||
numeric_fields = ("support", "precision", "recall", "f1")
|
||||
if task != "classify":
|
||||
numeric_fields += ("map50", "map50_95")
|
||||
for field in numeric_fields:
|
||||
try:
|
||||
value = float(row[field])
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise AssertionError(
|
||||
f"Invalid {field} in task_metrics.csv for {task}: {row}"
|
||||
) from exc
|
||||
require(
|
||||
math.isfinite(value),
|
||||
f"Non-finite {field} in task_metrics.csv for {task}: {row}",
|
||||
)
|
||||
return len(rows)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
|
|
@ -39,6 +225,10 @@ def main() -> None:
|
|||
default="sqlite:///runs/yolo26_mlflow_smoke/mlflow.db",
|
||||
)
|
||||
parser.add_argument("--experiment", default="yolo26-mlflow-smoke")
|
||||
parser.add_argument(
|
||||
"--run-group",
|
||||
help="Verify this yolo.run_group tag; defaults to the newest smoke batch.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
mlflow.set_tracking_uri(args.tracking_uri)
|
||||
|
|
@ -51,27 +241,75 @@ def main() -> None:
|
|||
[experiment.experiment_id],
|
||||
order_by=["start_time DESC"],
|
||||
)
|
||||
run_group = args.run_group
|
||||
if not run_group:
|
||||
run_group = latest_smoke_run_group(runs)
|
||||
|
||||
summary: dict[str, object] = {
|
||||
"tracking_uri": args.tracking_uri,
|
||||
"experiment_id": experiment.experiment_id,
|
||||
"artifact_location": experiment.artifact_location,
|
||||
"run_group": run_group,
|
||||
"tasks": {},
|
||||
}
|
||||
task_summary: dict[str, object] = summary["tasks"] # type: ignore[assignment]
|
||||
|
||||
for task in TASKS:
|
||||
run = latest_task_run(runs, task)
|
||||
run = latest_task_run(runs, task, run_group)
|
||||
artifacts = artifact_paths(client, run.info.run_id)
|
||||
missing = REQUIRED_ARTIFACTS - artifacts
|
||||
assert run.info.status == "FINISHED", (task, run.info.status)
|
||||
assert run.data.params, f"No parameters logged for {task}"
|
||||
assert run.data.metrics, f"No metrics logged for {task}"
|
||||
assert not missing, f"Missing artifacts for {task}: {sorted(missing)}"
|
||||
required_metrics = COMMON_MONITOR_METRICS | REQUIRED_METRICS_BY_TASK[task]
|
||||
missing_metrics = required_metrics - run.data.metrics.keys()
|
||||
require(
|
||||
run.info.status == "FINISHED",
|
||||
f"Unexpected run status for {task}: {run.info.status}",
|
||||
)
|
||||
require(run.data.params, f"No parameters logged for {task}")
|
||||
require(
|
||||
not missing_metrics,
|
||||
f"Missing metrics for {task}: {sorted(missing_metrics)}",
|
||||
)
|
||||
invalid_metrics = {
|
||||
key: run.data.metrics[key]
|
||||
for key in required_metrics
|
||||
if not math.isfinite(run.data.metrics[key])
|
||||
}
|
||||
require(
|
||||
not invalid_metrics,
|
||||
f"Non-finite metrics for {task}: {invalid_metrics}",
|
||||
)
|
||||
require(not missing, f"Missing artifacts for {task}: {sorted(missing)}")
|
||||
for metric_key in required_metrics:
|
||||
require(
|
||||
client.get_metric_history(run.info.run_id, metric_key),
|
||||
f"No metric history for {task}: {metric_key}",
|
||||
)
|
||||
history_key = HISTORY_METRIC_BY_TASK[task]
|
||||
history = client.get_metric_history(run.info.run_id, history_key)
|
||||
require(
|
||||
all(math.isfinite(point.value) for point in history),
|
||||
f"Non-finite metric history for {task}: {history_key} {history}",
|
||||
)
|
||||
history_steps = [point.step for point in history]
|
||||
require(
|
||||
history_steps == [0, 1],
|
||||
f"Unexpected metric steps for {task}: {history_key} {history_steps}",
|
||||
)
|
||||
require(run.data.tags.get("yolo.task") == task, f"Missing task tag for {task}")
|
||||
require(
|
||||
run.data.tags.get("monitoring.schema_version") == "1",
|
||||
f"Missing monitoring schema tag for {task}",
|
||||
)
|
||||
per_class_rows = verify_per_class_artifact(client, run.info.run_id, task)
|
||||
task_summary[task] = {
|
||||
"run_id": run.info.run_id,
|
||||
"status": run.info.status,
|
||||
"parameters": len(run.data.params),
|
||||
"metrics": len(run.data.metrics),
|
||||
"required_metrics": sorted(required_metrics),
|
||||
"history_metric": history_key,
|
||||
"history_steps": history_steps,
|
||||
"per_class_rows": per_class_rows,
|
||||
"artifact_uri": run.info.artifact_uri,
|
||||
"required_artifacts": sorted(REQUIRED_ARTIFACTS),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from fastapi.staticfiles import StaticFiles
|
|||
import uvicorn
|
||||
|
||||
from yolo_webui.config import TrainingConfig
|
||||
from yolo_webui.export_runner import ExportConfig
|
||||
from yolo_webui.trainer import TrainingRunner
|
||||
|
||||
# Set up logging
|
||||
|
|
@ -27,6 +28,17 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(me
|
|||
logger = logging.getLogger("yolo_webui")
|
||||
SESSION_NAME_PATTERN = re.compile(r"^[A-Za-z0-9_-]{1,64}$")
|
||||
|
||||
try:
|
||||
from ultralytics import settings as _ultralytics_settings
|
||||
_workspace_root = Path.cwd().resolve()
|
||||
_ultralytics_settings.update({
|
||||
"runs_dir": str((_workspace_root / "runs").resolve()),
|
||||
"datasets_dir": str((_workspace_root / "datasets").resolve()),
|
||||
"weights_dir": str((_workspace_root / "models").resolve()),
|
||||
})
|
||||
except Exception as _exc:
|
||||
logger.debug(f"Ultralytics settings init skipped: {_exc}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class LiveState:
|
||||
|
|
@ -49,6 +61,17 @@ class LiveState:
|
|||
self.stop_requested = False
|
||||
self.last_event_kind = None
|
||||
|
||||
def snapshot(self) -> dict[str, Any]:
|
||||
"""Return a detached snapshot safe to serialize after releasing a lock."""
|
||||
return {
|
||||
"status": self.status,
|
||||
"epoch": self.epoch,
|
||||
"total_epochs": self.total_epochs,
|
||||
"logs": list(self.logs),
|
||||
"metrics": [dict(metric) for metric in self.metrics],
|
||||
"output_dir": self.output_dir,
|
||||
}
|
||||
|
||||
|
||||
class TrainingManager:
|
||||
"""Manages the background training subprocess and WebSocket clients."""
|
||||
|
|
@ -355,6 +378,51 @@ def get_sessions_dir() -> Path:
|
|||
return path
|
||||
|
||||
|
||||
def write_json_atomic(path: Path, data: dict[str, Any]) -> None:
|
||||
"""Replace a JSON file atomically so concurrent readers never see partial data."""
|
||||
temp_path: Path | None = None
|
||||
try:
|
||||
with tempfile.NamedTemporaryFile(
|
||||
"w",
|
||||
prefix=f".{path.name}.",
|
||||
suffix=".tmp",
|
||||
dir=path.parent,
|
||||
delete=False,
|
||||
encoding="utf-8",
|
||||
) as temp_file:
|
||||
temp_path = Path(temp_file.name)
|
||||
json.dump(data, temp_file, ensure_ascii=False, indent=2)
|
||||
os.replace(temp_path, path)
|
||||
temp_path = None
|
||||
finally:
|
||||
if temp_path is not None:
|
||||
try:
|
||||
temp_path.unlink()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
|
||||
def parse_training_config(config_data: dict[str, Any]) -> TrainingConfig:
|
||||
config = TrainingConfig.from_dict(config_data)
|
||||
config.validate()
|
||||
return config
|
||||
|
||||
|
||||
def parse_export_config(config_data: dict[str, Any]) -> dict[str, Any]:
|
||||
config = ExportConfig.from_mapping(config_data)
|
||||
return {
|
||||
"model": str(config.model),
|
||||
"format": config.export_format,
|
||||
"imgsz": config.imgsz,
|
||||
"half": config.half,
|
||||
"int8": config.int8,
|
||||
"dynamic": config.dynamic,
|
||||
"simplify": config.simplify,
|
||||
"batch": config.batch,
|
||||
"workspace": config.workspace,
|
||||
}
|
||||
|
||||
|
||||
def get_session_path(name: str, *, allow_last_run: bool = True) -> Path:
|
||||
if SESSION_NAME_PATTERN.fullmatch(name) is None:
|
||||
raise HTTPException(
|
||||
|
|
@ -432,13 +500,13 @@ async def list_models():
|
|||
try:
|
||||
for path in runs_dir.rglob("*.pt"):
|
||||
items.append({
|
||||
"name": f"{path.parent.parent.parent.name}/{path.parent.parent.name}/{path.name}",
|
||||
"name": path.relative_to(runs_dir).as_posix(),
|
||||
"path": str(path.absolute()),
|
||||
"source": "runs"
|
||||
})
|
||||
for path in runs_dir.rglob("*.pth"):
|
||||
items.append({
|
||||
"name": f"{path.parent.parent.parent.name}/{path.parent.parent.name}/{path.name}",
|
||||
"name": path.relative_to(runs_dir).as_posix(),
|
||||
"path": str(path.absolute()),
|
||||
"source": "runs"
|
||||
})
|
||||
|
|
@ -464,8 +532,13 @@ async def load_session(name: str):
|
|||
async def save_session(name: str, config_data: dict[str, Any]):
|
||||
file_path = get_session_path(name, allow_last_run=False)
|
||||
try:
|
||||
with file_path.open("w", encoding="utf-8") as f:
|
||||
json.dump(config_data, f, ensure_ascii=False, indent=2)
|
||||
config = parse_training_config(config_data)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc))
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=400, detail=f"Некорректная конфигурация: {exc}")
|
||||
try:
|
||||
write_json_atomic(file_path, config.to_dict())
|
||||
return {"message": "Сессия успешно сохранена."}
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=500, detail=f"Не удалось сохранить сессию: {exc}")
|
||||
|
|
@ -486,40 +559,32 @@ async def delete_session(name: str):
|
|||
@app.get("/api/train/status")
|
||||
async def get_status():
|
||||
with manager._lock:
|
||||
return {
|
||||
"status": manager.state.status,
|
||||
"epoch": manager.state.epoch,
|
||||
"total_epochs": manager.state.total_epochs,
|
||||
"output_dir": manager.state.output_dir,
|
||||
"metrics": manager.state.metrics,
|
||||
"logs": manager.state.logs,
|
||||
}
|
||||
return manager.state.snapshot()
|
||||
|
||||
|
||||
@app.post("/api/train/start")
|
||||
async def start_training(config_data: dict[str, Any]):
|
||||
try:
|
||||
config = TrainingConfig.from_dict(config_data)
|
||||
config.validate()
|
||||
config = parse_training_config(config_data)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc))
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=400, detail=f"Некорректная конфигурация: {exc}")
|
||||
|
||||
# Auto-save last configuration on start
|
||||
try:
|
||||
manager.start_training(config)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc))
|
||||
|
||||
# Only a configuration accepted for execution becomes the last run.
|
||||
try:
|
||||
sessions_dir = get_sessions_dir()
|
||||
last_run_path = sessions_dir / "last_run.json"
|
||||
with last_run_path.open("w", encoding="utf-8") as f:
|
||||
json.dump(config_data, f, ensure_ascii=False, indent=2)
|
||||
write_json_atomic(last_run_path, config.to_dict())
|
||||
except Exception as exc:
|
||||
logger.error(f"Failed to auto-save last run: {exc}")
|
||||
|
||||
try:
|
||||
manager.start_training(config)
|
||||
return {"message": "Обучение запущено."}
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc))
|
||||
return {"message": "Обучение запущено."}
|
||||
|
||||
|
||||
@app.post("/api/train/stop")
|
||||
|
|
@ -572,8 +637,10 @@ class ExportManager:
|
|||
future = loop.create_task(coroutine)
|
||||
else:
|
||||
future = asyncio.run_coroutine_threadsafe(coroutine, loop)
|
||||
future.add_done_callback(TrainingManager._log_broadcast_failure)
|
||||
except Exception:
|
||||
coroutine.close()
|
||||
logger.exception("Failed to schedule export WebSocket broadcast")
|
||||
|
||||
async def _send_payload(self, payload: str) -> None:
|
||||
broadcast_lock = self._broadcast_lock
|
||||
|
|
@ -621,19 +688,28 @@ class ExportManager:
|
|||
with self._lock:
|
||||
if self.state.status not in ("preparing", "exporting"):
|
||||
return
|
||||
|
||||
process = self._process
|
||||
if process is not None:
|
||||
if process.poll() is not None:
|
||||
# The subprocess has already completed; let its real return code
|
||||
# determine the final status instead of relabeling it cancelled.
|
||||
return
|
||||
try:
|
||||
process.terminate()
|
||||
except Exception:
|
||||
logger.warning("Failed to terminate export subprocess", exc_info=True)
|
||||
return
|
||||
|
||||
self.state.status = "stopping"
|
||||
self.state.stop_requested = True
|
||||
if self._process is not None:
|
||||
try:
|
||||
self._process.terminate()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
self.broadcast({"type": "status", "status": self.state.status})
|
||||
self.add_log("Запрошена остановка экспорта...", "warning")
|
||||
|
||||
def _run_subprocess(self, config_data: dict[str, Any]) -> None:
|
||||
temp_config_path = None
|
||||
process: subprocess.Popen[str] | None = None
|
||||
try:
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(config_data, f)
|
||||
|
|
@ -641,17 +717,29 @@ class ExportManager:
|
|||
|
||||
cmd = [sys.executable, "-u", "-m", "yolo_webui.export_runner", temp_config_path]
|
||||
with self._lock:
|
||||
self._process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
if self.state.stop_requested:
|
||||
cancelled_before_start = True
|
||||
else:
|
||||
cancelled_before_start = False
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
self._process = process
|
||||
|
||||
if cancelled_before_start:
|
||||
self._finalize_process_result(0)
|
||||
return
|
||||
|
||||
if process is None or process.stdout is None:
|
||||
raise RuntimeError("Не удалось открыть stdout процесса экспорта.")
|
||||
|
||||
buffer = ""
|
||||
while True:
|
||||
char = self._process.stdout.read(1)
|
||||
char = process.stdout.read(1)
|
||||
if not char:
|
||||
if buffer:
|
||||
self._handle_subprocess_line(buffer)
|
||||
|
|
@ -664,26 +752,28 @@ class ExportManager:
|
|||
else:
|
||||
buffer += char
|
||||
|
||||
self._process.wait()
|
||||
rc = self._process.returncode
|
||||
process.wait()
|
||||
rc = process.returncode
|
||||
|
||||
with self._lock:
|
||||
self._process = None
|
||||
if self._process is process:
|
||||
self._process = None
|
||||
|
||||
self._finalize_process_result(rc)
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception("Error in export process thread:")
|
||||
if self._process is not None:
|
||||
if process is not None:
|
||||
try:
|
||||
if self._process.poll() is None:
|
||||
self._process.kill()
|
||||
self._process.wait(timeout=5)
|
||||
if process.poll() is None:
|
||||
process.kill()
|
||||
process.wait(timeout=5)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
with self._lock:
|
||||
self._process = None
|
||||
if self._process is process:
|
||||
self._process = None
|
||||
self.state.status = "failed"
|
||||
self.add_log(f"Внутренняя ошибка менеджера: {exc}", "error")
|
||||
finally:
|
||||
|
|
@ -703,9 +793,13 @@ class ExportManager:
|
|||
return
|
||||
|
||||
if line_str == "__YOLO_WEBUI_READY__":
|
||||
status_changed = False
|
||||
with self._lock:
|
||||
self.state.status = "exporting"
|
||||
self.broadcast({"type": "status", "status": "exporting"})
|
||||
if self.state.status == "preparing":
|
||||
self.state.status = "exporting"
|
||||
status_changed = True
|
||||
if status_changed:
|
||||
self.broadcast({"type": "status", "status": "exporting"})
|
||||
elif line_str.startswith("__YOLO_WEBUI_RESULT__:"):
|
||||
with self._lock:
|
||||
self.state.output_dir = line_str[len("__YOLO_WEBUI_RESULT__:") :]
|
||||
|
|
@ -738,16 +832,18 @@ export_manager = ExportManager()
|
|||
@app.get("/api/export/status")
|
||||
async def get_export_status():
|
||||
with export_manager._lock:
|
||||
snapshot = export_manager.state.snapshot()
|
||||
return {
|
||||
"status": export_manager.state.status,
|
||||
"output_dir": export_manager.state.output_dir,
|
||||
"logs": export_manager.state.logs,
|
||||
"status": snapshot["status"],
|
||||
"output_dir": snapshot["output_dir"],
|
||||
"logs": snapshot["logs"],
|
||||
}
|
||||
|
||||
@app.post("/api/export/start")
|
||||
async def start_export(config_data: dict[str, Any]):
|
||||
try:
|
||||
export_manager.start_export(config_data)
|
||||
validated_data = parse_export_config(config_data)
|
||||
export_manager.start_export(validated_data)
|
||||
return {"message": "Экспорт запущен."}
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc))
|
||||
|
|
@ -763,11 +859,12 @@ async def export_websocket_endpoint(websocket: WebSocket):
|
|||
export_manager.add_websocket(websocket)
|
||||
|
||||
with export_manager._lock:
|
||||
snapshot = export_manager.state.snapshot()
|
||||
state_dict = {
|
||||
"type": "init",
|
||||
"status": export_manager.state.status,
|
||||
"output_dir": export_manager.state.output_dir,
|
||||
"logs": [log.split(":", 1) for log in export_manager.state.logs if ":" in log],
|
||||
"status": snapshot["status"],
|
||||
"output_dir": snapshot["output_dir"],
|
||||
"logs": [log.split(":", 1) for log in snapshot["logs"] if ":" in log],
|
||||
}
|
||||
try:
|
||||
await websocket.send_text(json.dumps(state_dict))
|
||||
|
|
@ -788,15 +885,16 @@ async def websocket_endpoint(websocket: WebSocket):
|
|||
|
||||
# Send current state upon connection
|
||||
with manager._lock:
|
||||
snapshot = manager.state.snapshot()
|
||||
state_dict = {
|
||||
"type": "init",
|
||||
"status": manager.state.status,
|
||||
"epoch": manager.state.epoch,
|
||||
"total_epochs": manager.state.total_epochs,
|
||||
"output_dir": manager.state.output_dir,
|
||||
"metrics": manager.state.metrics,
|
||||
"status": snapshot["status"],
|
||||
"epoch": snapshot["epoch"],
|
||||
"total_epochs": snapshot["total_epochs"],
|
||||
"output_dir": snapshot["output_dir"],
|
||||
"metrics": snapshot["metrics"],
|
||||
# We format log items for the UI
|
||||
"logs": [log.split(":", 1) for log in manager.state.logs if ":" in log],
|
||||
"logs": [log.split(":", 1) for log in snapshot["logs"] if ":" in log],
|
||||
}
|
||||
try:
|
||||
await websocket.send_text(json.dumps(state_dict))
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
import math
|
||||
import os
|
||||
from pathlib import Path
|
||||
from pathlib import Path, PureWindowsPath
|
||||
import re
|
||||
from typing import Any, Literal
|
||||
|
||||
|
|
@ -28,6 +29,58 @@ SAFE_IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*$")
|
|||
MODEL_SUFFIXES = {".pt", ".pth", ".yaml", ".yml"}
|
||||
|
||||
|
||||
def _default_mlflow_tracking_uri() -> str:
|
||||
return os.environ.get(
|
||||
"YOLO_WEBUI_MLFLOW_TRACKING_URI",
|
||||
"sqlite:///mlflow.db",
|
||||
)
|
||||
|
||||
|
||||
def _require_string(value: Any, label: str) -> str:
|
||||
if not isinstance(value, str):
|
||||
raise ValueError(f"{label} должен быть строкой.")
|
||||
return value
|
||||
|
||||
|
||||
def _require_bool(value: Any, label: str) -> bool:
|
||||
if not isinstance(value, bool):
|
||||
raise ValueError(f"{label} должен быть логическим значением.")
|
||||
return value
|
||||
|
||||
|
||||
def _require_finite_number(value: Any, label: str) -> float:
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
raise ValueError(f"{label} должен быть числом.")
|
||||
try:
|
||||
number = float(value)
|
||||
except OverflowError as exc:
|
||||
raise ValueError(f"{label} должен быть конечным числом.") from exc
|
||||
if not math.isfinite(number):
|
||||
raise ValueError(f"{label} должен быть конечным числом.")
|
||||
return number
|
||||
|
||||
|
||||
def _require_integer(value: Any, label: str) -> int:
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise ValueError(f"{label} должен быть целым числом.")
|
||||
return value
|
||||
|
||||
|
||||
def _validate_run_name(value: str) -> None:
|
||||
normalized = value.strip()
|
||||
if not normalized:
|
||||
return
|
||||
if (
|
||||
normalized in {".", ".."}
|
||||
or "/" in normalized
|
||||
or "\\" in normalized
|
||||
or PureWindowsPath(normalized).drive
|
||||
):
|
||||
raise ValueError(
|
||||
"Имя запуска должно быть одним именем каталога без '/' или '\\'."
|
||||
)
|
||||
|
||||
|
||||
def _allowed_roots(defaults: tuple[str, ...], environment_name: str) -> tuple[Path, ...]:
|
||||
configured = [
|
||||
item
|
||||
|
|
@ -109,7 +162,7 @@ class AugmentationConfig:
|
|||
close_mosaic: int = 10
|
||||
|
||||
def validate(self) -> None:
|
||||
if not self.enabled:
|
||||
if not _require_bool(self.enabled, "Флаг аугментации"):
|
||||
return
|
||||
|
||||
fractions = {
|
||||
|
|
@ -129,20 +182,23 @@ class AugmentationConfig:
|
|||
"Erasing": self.erasing,
|
||||
}
|
||||
for label, value in fractions.items():
|
||||
if not 0.0 <= value <= 1.0:
|
||||
number = _require_finite_number(value, f"Параметр «{label}»")
|
||||
if not 0.0 <= number <= 1.0:
|
||||
raise ValueError(f"Параметр «{label}» должен быть от 0 до 1.")
|
||||
if self.degrees < 0:
|
||||
if _require_finite_number(self.degrees, "Угол поворота") < 0:
|
||||
raise ValueError("Угол поворота не может быть отрицательным.")
|
||||
if self.shear < 0:
|
||||
if _require_finite_number(self.shear, "Угол сдвига") < 0:
|
||||
raise ValueError("Угол сдвига не может быть отрицательным.")
|
||||
if self.close_mosaic < 0:
|
||||
if _require_integer(self.close_mosaic, "Close mosaic") < 0:
|
||||
raise ValueError("Close mosaic не может быть отрицательным.")
|
||||
_require_string(self.copy_paste_mode, "Режим copy-paste")
|
||||
if self.copy_paste_mode not in SUPPORTED_COPY_PASTE_MODES:
|
||||
raise ValueError(f"Неизвестный режим copy-paste: {self.copy_paste_mode}.")
|
||||
_require_string(self.auto_augment, "Политика AutoAugment")
|
||||
if self.auto_augment not in SUPPORTED_AUTO_AUGMENT_POLICIES:
|
||||
raise ValueError(f"Неизвестная политика AutoAugment: {self.auto_augment}.")
|
||||
|
||||
def train_kwargs(self) -> dict[str, str | int | float]:
|
||||
def train_kwargs(self) -> dict[str, str | int | float | None]:
|
||||
if not self.enabled:
|
||||
return {}
|
||||
return {
|
||||
|
|
@ -171,11 +227,15 @@ class AugmentationConfig:
|
|||
@dataclass(frozen=True, slots=True)
|
||||
class MlflowConfig:
|
||||
enabled: bool = True
|
||||
tracking_uri: str = "sqlite:///mlflow.db"
|
||||
tracking_uri: str = field(default_factory=_default_mlflow_tracking_uri)
|
||||
experiment_name: str = "yolo-webui"
|
||||
run_name: str = ""
|
||||
|
||||
def validate(self) -> None:
|
||||
_require_bool(self.enabled, "Флаг MLflow")
|
||||
_require_string(self.tracking_uri, "URI хранилища MLflow")
|
||||
_require_string(self.experiment_name, "Название эксперимента MLflow")
|
||||
_require_string(self.run_name, "Название запуска MLflow")
|
||||
if self.enabled and not self.tracking_uri.strip():
|
||||
raise ValueError("Укажите URI хранилища MLflow.")
|
||||
if self.enabled and not self.experiment_name.strip():
|
||||
|
|
@ -189,8 +249,14 @@ class DatasetSplitConfig:
|
|||
classes_path: str = ""
|
||||
|
||||
def validate(self) -> None:
|
||||
_require_bool(self.enabled, "Флаг разделения датасета")
|
||||
ratio = _require_finite_number(
|
||||
self.train_ratio,
|
||||
"Доля обучающей выборки (Train)",
|
||||
)
|
||||
_require_string(self.classes_path, "Путь к файлу классов")
|
||||
if self.enabled:
|
||||
if not 0.1 <= self.train_ratio <= 0.95:
|
||||
if not 0.1 <= ratio <= 0.95:
|
||||
raise ValueError("Доля обучающей выборки (Train) должна быть от 0.1 до 0.95.")
|
||||
|
||||
|
||||
|
|
@ -212,10 +278,27 @@ class TrainingConfig:
|
|||
split: DatasetSplitConfig = field(default_factory=DatasetSplitConfig)
|
||||
|
||||
def validate(self) -> None:
|
||||
if not self.dataset.strip():
|
||||
dataset = _require_string(self.dataset, "Датасет")
|
||||
model = _require_string(self.model, "Модель")
|
||||
_require_string(self.task, "Тип задачи")
|
||||
_require_string(self.device, "Устройство")
|
||||
project = _require_string(self.project, "Каталог результатов")
|
||||
run_name = _require_string(self.run_name, "Имя запуска")
|
||||
if not isinstance(self.augmentation, AugmentationConfig):
|
||||
raise ValueError("Параметры аугментации должны быть объектом.")
|
||||
if not isinstance(self.mlflow, MlflowConfig):
|
||||
raise ValueError("Параметры MLflow должны быть объектом.")
|
||||
if not isinstance(self.split, DatasetSplitConfig):
|
||||
raise ValueError("Параметры разделения датасета должны быть объектом.")
|
||||
self.augmentation.validate()
|
||||
self.mlflow.validate()
|
||||
self.split.validate()
|
||||
|
||||
if not dataset.strip():
|
||||
raise ValueError("Укажите путь или имя датасета.")
|
||||
if not self.model.strip():
|
||||
if not model.strip():
|
||||
raise ValueError("Укажите путь или имя модели.")
|
||||
_validate_run_name(run_name)
|
||||
|
||||
data_roots = _allowed_roots(("datasets",), "YOLO_WEBUI_DATA_ROOTS")
|
||||
model_roots = _allowed_roots(
|
||||
|
|
@ -224,23 +307,29 @@ class TrainingConfig:
|
|||
)
|
||||
run_roots = _allowed_roots(("runs",), "YOLO_WEBUI_RUN_ROOTS")
|
||||
_validate_local_reference(
|
||||
self.dataset,
|
||||
dataset,
|
||||
label="Датасет",
|
||||
roots=data_roots,
|
||||
allow_identifier=True,
|
||||
)
|
||||
_validate_local_reference(
|
||||
self.model,
|
||||
model,
|
||||
label="Модель",
|
||||
roots=model_roots,
|
||||
allow_identifier=True,
|
||||
allowed_suffixes=MODEL_SUFFIXES,
|
||||
)
|
||||
_validate_local_reference(
|
||||
self.project.strip() or "runs/train",
|
||||
project.strip() or "runs/train",
|
||||
label="Каталог результатов",
|
||||
roots=run_roots,
|
||||
)
|
||||
if run_name.strip():
|
||||
_validate_local_reference(
|
||||
str(Path(project.strip() or "runs/train") / run_name.strip()),
|
||||
label="Каталог запуска",
|
||||
roots=run_roots,
|
||||
)
|
||||
if self.split.classes_path.strip():
|
||||
_validate_local_reference(
|
||||
self.split.classes_path,
|
||||
|
|
@ -254,22 +343,31 @@ class TrainingConfig:
|
|||
"Автоматическое разделение доступно только для YOLO-датасетов "
|
||||
"с папками images/labels и не поддерживает задачу classify."
|
||||
)
|
||||
if self.epochs < 1:
|
||||
if _require_integer(self.epochs, "Количество эпох") < 1:
|
||||
raise ValueError("Количество эпох должно быть не меньше 1.")
|
||||
if self.image_size < 32:
|
||||
if _require_integer(self.image_size, "Размер изображения") < 32:
|
||||
raise ValueError("Размер изображения должен быть не меньше 32.")
|
||||
if self.batch_size == 0 or self.batch_size < -1:
|
||||
raise ValueError("Batch должен быть положительным числом или -1 для автоподбора.")
|
||||
if self.workers < 0:
|
||||
batch_size = _require_finite_number(self.batch_size, "Batch")
|
||||
if isinstance(self.batch_size, int):
|
||||
valid_batch = self.batch_size == -1 or self.batch_size >= 1
|
||||
else:
|
||||
valid_batch = 0 < batch_size < 1
|
||||
if not valid_batch:
|
||||
raise ValueError(
|
||||
"Batch должен быть целым числом от 1, значением -1 для "
|
||||
"автоподбора либо дробью от 0 до 1."
|
||||
)
|
||||
if _require_integer(self.workers, "Количество workers") < 0:
|
||||
raise ValueError("Количество workers не может быть отрицательным.")
|
||||
if self.patience < 0:
|
||||
if _require_integer(self.patience, "Patience") < 0:
|
||||
raise ValueError("Patience не может быть отрицательным.")
|
||||
self.augmentation.validate()
|
||||
self.mlflow.validate()
|
||||
self.split.validate()
|
||||
|
||||
def train_kwargs(self) -> dict[str, str | int | float | bool]:
|
||||
"""Convert the form values to arguments accepted by YOLO.train()."""
|
||||
project_path = Path(self.project.strip() or "runs/train")
|
||||
if not project_path.is_absolute():
|
||||
project_path = (Path.cwd() / project_path).resolve()
|
||||
|
||||
values: dict[str, str | int | float | bool] = {
|
||||
"data": self.dataset.strip(),
|
||||
"epochs": self.epochs,
|
||||
|
|
@ -277,7 +375,7 @@ class TrainingConfig:
|
|||
"batch": self.batch_size,
|
||||
"workers": self.workers,
|
||||
"patience": self.patience,
|
||||
"project": self.project.strip() or "runs/train",
|
||||
"project": str(project_path),
|
||||
# Enable verbose output so users see active progress and losses in the log console.
|
||||
"verbose": True,
|
||||
}
|
||||
|
|
@ -303,12 +401,30 @@ class TrainingConfig:
|
|||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> TrainingConfig:
|
||||
aug_data = data.get("augmentation", {})
|
||||
mlflow_data = data.get("mlflow", {})
|
||||
split_data = data.get("split", {})
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("Конфигурация должна быть JSON-объектом.")
|
||||
|
||||
sections: dict[str, dict[str, Any]] = {}
|
||||
for key, label in (
|
||||
("augmentation", "аугментации"),
|
||||
("mlflow", "MLflow"),
|
||||
("split", "разделения датасета"),
|
||||
):
|
||||
raw_section = data.get(key, {})
|
||||
if not isinstance(raw_section, dict):
|
||||
raise ValueError(f"Параметры {label} должны быть JSON-объектом.")
|
||||
sections[key] = raw_section
|
||||
|
||||
try:
|
||||
augmentation = AugmentationConfig(**sections["augmentation"])
|
||||
mlflow = MlflowConfig(**sections["mlflow"])
|
||||
split = DatasetSplitConfig(**sections["split"])
|
||||
except TypeError as exc:
|
||||
raise ValueError(f"Конфигурация содержит неизвестные параметры: {exc}") from exc
|
||||
|
||||
return cls(
|
||||
dataset=data["dataset"],
|
||||
model=data["model"],
|
||||
dataset=data.get("dataset", ""),
|
||||
model=data.get("model", ""),
|
||||
task=data.get("task", "detect"),
|
||||
epochs=data.get("epochs", 100),
|
||||
image_size=data.get("image_size", 640),
|
||||
|
|
@ -318,7 +434,7 @@ class TrainingConfig:
|
|||
patience=data.get("patience", 100),
|
||||
project=data.get("project", "runs/train"),
|
||||
run_name=data.get("run_name", ""),
|
||||
augmentation=AugmentationConfig(**aug_data) if aug_data else AugmentationConfig(),
|
||||
mlflow=MlflowConfig(**mlflow_data) if mlflow_data else MlflowConfig(),
|
||||
split=DatasetSplitConfig(**split_data) if split_data else DatasetSplitConfig(),
|
||||
augmentation=augmentation,
|
||||
mlflow=mlflow,
|
||||
split=split,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
|
@ -8,6 +9,25 @@ from uuid import uuid4
|
|||
import yaml
|
||||
|
||||
|
||||
def _normalize_class_key(raw_key: Any, source: Path) -> int:
|
||||
if isinstance(raw_key, bool):
|
||||
raise ValueError(f"Некорректный ID класса в '{source}': {raw_key!r}.")
|
||||
if isinstance(raw_key, int):
|
||||
class_id = raw_key
|
||||
elif isinstance(raw_key, str):
|
||||
try:
|
||||
class_id = int(raw_key.strip())
|
||||
except ValueError as exc:
|
||||
raise ValueError(
|
||||
f"Некорректный ID класса в '{source}': {raw_key!r}."
|
||||
) from exc
|
||||
else:
|
||||
raise ValueError(f"Некорректный ID класса в '{source}': {raw_key!r}.")
|
||||
if class_id < 0:
|
||||
raise ValueError(f"ID класса в '{source}' не может быть отрицательным.")
|
||||
return class_id
|
||||
|
||||
|
||||
def _normalize_names(names: Any, source: Path) -> dict[int, str]:
|
||||
if isinstance(names, list):
|
||||
items = enumerate(names)
|
||||
|
|
@ -15,14 +35,7 @@ def _normalize_names(names: Any, source: Path) -> dict[int, str]:
|
|||
normalized_items: list[tuple[int, Any]] = []
|
||||
seen: set[int] = set()
|
||||
for raw_key, value in names.items():
|
||||
if isinstance(raw_key, bool):
|
||||
raise ValueError(f"Некорректный ID класса в '{source}': {raw_key!r}.")
|
||||
try:
|
||||
class_id = int(raw_key)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(
|
||||
f"Некорректный ID класса в '{source}': {raw_key!r}."
|
||||
) from exc
|
||||
class_id = _normalize_class_key(raw_key, source)
|
||||
if class_id in seen:
|
||||
raise ValueError(f"Повторяющийся ID класса {class_id} в '{source}'.")
|
||||
seen.add(class_id)
|
||||
|
|
@ -86,6 +99,36 @@ def _parse_classes_file(path: Path) -> dict[int, str]:
|
|||
)
|
||||
|
||||
|
||||
def _read_label_class_ids(path: Path) -> set[int]:
|
||||
class_ids: set[int] = set()
|
||||
try:
|
||||
with path.open("r", encoding="utf-8") as label_file:
|
||||
for line_number, line in enumerate(label_file, start=1):
|
||||
parts = line.strip().split()
|
||||
if not parts:
|
||||
continue
|
||||
try:
|
||||
numeric_id = float(parts[0])
|
||||
except ValueError as exc:
|
||||
raise ValueError(
|
||||
f"Некорректный ID класса в '{path}', строка {line_number}: "
|
||||
f"{parts[0]!r}."
|
||||
) from exc
|
||||
if (
|
||||
not math.isfinite(numeric_id)
|
||||
or not numeric_id.is_integer()
|
||||
or numeric_id < 0
|
||||
):
|
||||
raise ValueError(
|
||||
f"Некорректный ID класса в '{path}', строка {line_number}: "
|
||||
f"{parts[0]!r}."
|
||||
)
|
||||
class_ids.add(int(numeric_id))
|
||||
except (OSError, UnicodeError) as exc:
|
||||
raise ValueError(f"Не удалось прочитать файл разметки '{path}': {exc}.") from exc
|
||||
return class_ids
|
||||
|
||||
|
||||
def read_classes(dataset_dir: Path, custom_classes_path: str) -> dict[int, str]:
|
||||
# An explicit path is authoritative: typos and malformed files must not fall back.
|
||||
if custom_classes_path.strip():
|
||||
|
|
@ -119,17 +162,10 @@ def read_classes(dataset_dir: Path, custom_classes_path: str) -> dict[int, str]:
|
|||
class_ids: set[int] = set()
|
||||
labels_dir = dataset_dir / "labels"
|
||||
if labels_dir.exists():
|
||||
for txt_file in labels_dir.rglob("*.txt"):
|
||||
for txt_file in sorted(labels_dir.rglob("*.txt")):
|
||||
if txt_file.name == "classes.txt":
|
||||
continue
|
||||
try:
|
||||
with txt_file.open("r", encoding="utf-8") as label_file:
|
||||
for line in label_file:
|
||||
parts = line.strip().split()
|
||||
if parts:
|
||||
class_ids.add(int(parts[0]))
|
||||
except (OSError, ValueError):
|
||||
continue
|
||||
class_ids.update(_read_label_class_ids(txt_file))
|
||||
|
||||
if class_ids:
|
||||
max_id = max(class_ids)
|
||||
|
|
@ -146,15 +182,52 @@ def _write_new(path: Path, content: str) -> None:
|
|||
output_file.write(content)
|
||||
|
||||
|
||||
def _validate_image_labels(
|
||||
images_dir: Path,
|
||||
labels_dir: Path,
|
||||
image_files: list[Path],
|
||||
classes: dict[int, str],
|
||||
) -> None:
|
||||
checked_paths: set[Path] = set()
|
||||
for image_path in image_files:
|
||||
label_path = labels_dir / image_path.relative_to(images_dir).with_suffix(".txt")
|
||||
if label_path in checked_paths:
|
||||
continue
|
||||
checked_paths.add(label_path)
|
||||
if not label_path.exists():
|
||||
# YOLO treats an image without a label file as a background image.
|
||||
continue
|
||||
if not label_path.is_file():
|
||||
raise ValueError(f"Путь к разметке '{label_path}' не является файлом.")
|
||||
for class_id in _read_label_class_ids(label_path):
|
||||
if class_id not in classes:
|
||||
raise ValueError(
|
||||
f"В файле '{label_path}' указан класс {class_id}, "
|
||||
f"но список классов содержит ID от 0 до {len(classes) - 1}."
|
||||
)
|
||||
|
||||
|
||||
def split_dataset(
|
||||
dataset_dir: str, train_ratio: float, classes_path: str
|
||||
) -> tuple[int, int, str]:
|
||||
if not isinstance(dataset_dir, str):
|
||||
raise ValueError("Каталог датасета должен быть строкой.")
|
||||
if not dataset_dir.strip():
|
||||
raise ValueError("Укажите каталог датасета.")
|
||||
if not 0.1 <= train_ratio <= 0.95:
|
||||
if not isinstance(classes_path, str):
|
||||
raise ValueError("Путь к файлу классов должен быть строкой.")
|
||||
if isinstance(train_ratio, bool) or not isinstance(train_ratio, (int, float)):
|
||||
raise ValueError("Доля обучающей выборки должна быть числом.")
|
||||
try:
|
||||
ratio = float(train_ratio)
|
||||
except OverflowError as exc:
|
||||
raise ValueError("Доля обучающей выборки должна быть конечным числом.") from exc
|
||||
if not math.isfinite(ratio):
|
||||
raise ValueError("Доля обучающей выборки должна быть конечным числом.")
|
||||
if not 0.1 <= ratio <= 0.95:
|
||||
raise ValueError("Доля обучающей выборки должна быть от 0.1 до 0.95.")
|
||||
|
||||
base_dir = Path(dataset_dir.strip()).expanduser().absolute()
|
||||
base_dir = Path(dataset_dir.strip()).expanduser().resolve(strict=False)
|
||||
images_dir = base_dir / "images"
|
||||
labels_dir = base_dir / "labels"
|
||||
|
||||
|
|
@ -185,13 +258,14 @@ def split_dataset(
|
|||
rng = random.Random(42)
|
||||
rng.shuffle(image_files)
|
||||
|
||||
split_idx = int(len(image_files) * train_ratio)
|
||||
split_idx = int(len(image_files) * ratio)
|
||||
split_idx = max(1, min(split_idx, len(image_files) - 1))
|
||||
train_images = image_files[:split_idx]
|
||||
val_images = image_files[split_idx:]
|
||||
|
||||
# Resolve classes before creating output so invalid input leaves no partial split.
|
||||
classes = read_classes(base_dir, classes_path)
|
||||
_validate_image_labels(images_dir, labels_dir, image_files, classes)
|
||||
|
||||
relative_split_dir = Path(".yolo-webui") / "splits" / uuid4().hex
|
||||
split_dir = base_dir / relative_split_dir
|
||||
|
|
|
|||
|
|
@ -1,14 +1,213 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
# Force headless Matplotlib to avoid any thread/process GUI issues
|
||||
from yolo_webui.config import MODEL_SUFFIXES, _allowed_roots, _is_within
|
||||
|
||||
# Force headless Matplotlib to avoid any thread/process GUI issues.
|
||||
os.environ["MPLBACKEND"] = "Agg"
|
||||
# Restrict PyTorch checkpoint deserialization for direct CLI invocation too.
|
||||
os.environ["ULTRALYTICS_SAFE_LOAD"] = "1"
|
||||
|
||||
def main() -> int:
|
||||
args = sys.argv[1:]
|
||||
SUPPORTED_EXPORT_FORMATS = frozenset(
|
||||
{
|
||||
"coreml",
|
||||
"engine",
|
||||
"onnx",
|
||||
"openvino",
|
||||
"pb",
|
||||
"saved_model",
|
||||
"tflite",
|
||||
"torchscript",
|
||||
}
|
||||
)
|
||||
EXPORT_CONFIG_FIELDS = frozenset(
|
||||
{
|
||||
"model",
|
||||
"format",
|
||||
"imgsz",
|
||||
"half",
|
||||
"int8",
|
||||
"dynamic",
|
||||
"simplify",
|
||||
"batch",
|
||||
"workspace",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _bounded_int(
|
||||
data: Mapping[str, Any],
|
||||
key: str,
|
||||
default: int,
|
||||
*,
|
||||
minimum: int,
|
||||
maximum: int,
|
||||
) -> int:
|
||||
value = data.get(key, default)
|
||||
if type(value) is not int or not minimum <= value <= maximum:
|
||||
raise ValueError(
|
||||
f"Параметр {key} должен быть целым числом от {minimum} до {maximum}."
|
||||
)
|
||||
return value
|
||||
|
||||
|
||||
def _boolean(data: Mapping[str, Any], key: str, default: bool = False) -> bool:
|
||||
value = data.get(key, default)
|
||||
if type(value) is not bool:
|
||||
raise ValueError(f"Параметр {key} должен быть логическим значением.")
|
||||
return value
|
||||
|
||||
|
||||
def _bounded_number(
|
||||
data: Mapping[str, Any],
|
||||
key: str,
|
||||
default: float,
|
||||
*,
|
||||
minimum: float,
|
||||
maximum: float,
|
||||
) -> float:
|
||||
value = data.get(key, default)
|
||||
if type(value) not in (int, float):
|
||||
raise ValueError(
|
||||
f"Параметр {key} должен быть числом от {minimum:g} до {maximum:g}."
|
||||
)
|
||||
try:
|
||||
number = float(value)
|
||||
except OverflowError as exc:
|
||||
raise ValueError(
|
||||
f"Параметр {key} должен быть числом от {minimum:g} до {maximum:g}."
|
||||
) from exc
|
||||
if not math.isfinite(number) or not minimum <= number <= maximum:
|
||||
raise ValueError(
|
||||
f"Параметр {key} должен быть числом от {minimum:g} до {maximum:g}."
|
||||
)
|
||||
return number
|
||||
|
||||
|
||||
def _model_path(value: Any) -> Path:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise ValueError("Укажите путь к модели.")
|
||||
|
||||
normalized = value.strip()
|
||||
if any(character in normalized for character in ("\0", "\r", "\n")):
|
||||
raise ValueError("Путь к модели содержит недопустимые управляющие символы.")
|
||||
|
||||
candidate = Path(normalized).expanduser()
|
||||
if not candidate.is_absolute():
|
||||
if len(candidate.parts) == 1:
|
||||
candidate = Path("models") / candidate
|
||||
candidate = Path.cwd() / candidate
|
||||
|
||||
try:
|
||||
resolved = candidate.resolve(strict=False)
|
||||
except (OSError, RuntimeError, ValueError) as exc:
|
||||
raise ValueError(f"Некорректный путь к модели: {normalized}") from exc
|
||||
|
||||
roots = _allowed_roots(("models", "runs"), "YOLO_WEBUI_MODEL_ROOTS")
|
||||
if not _is_within(resolved, roots):
|
||||
allowed = ", ".join(str(root) for root in roots)
|
||||
raise ValueError(
|
||||
f"Модель должна находиться в разрешённом каталоге: {allowed}."
|
||||
)
|
||||
try:
|
||||
is_file = resolved.is_file()
|
||||
except OSError as exc:
|
||||
raise ValueError(f"Не удалось проверить файл модели: {normalized}") from exc
|
||||
if not is_file:
|
||||
raise ValueError(f"Файл модели не найден: {normalized}")
|
||||
if resolved.suffix.lower() not in MODEL_SUFFIXES:
|
||||
expected = ", ".join(sorted(MODEL_SUFFIXES))
|
||||
raise ValueError(f"Модель должна иметь расширение {expected}.")
|
||||
return resolved
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ExportConfig:
|
||||
model: Path
|
||||
export_format: str
|
||||
imgsz: int
|
||||
half: bool
|
||||
int8: bool
|
||||
dynamic: bool
|
||||
simplify: bool
|
||||
batch: int
|
||||
workspace: float
|
||||
|
||||
@classmethod
|
||||
def from_mapping(cls, data: Mapping[str, Any]) -> ExportConfig:
|
||||
unknown = set(data) - EXPORT_CONFIG_FIELDS
|
||||
if unknown:
|
||||
fields = ", ".join(sorted(map(str, unknown)))
|
||||
raise ValueError(f"Неизвестные параметры экспорта: {fields}.")
|
||||
|
||||
raw_format = data.get("format", "onnx")
|
||||
if not isinstance(raw_format, str):
|
||||
raise ValueError("Параметр format должен быть строкой.")
|
||||
export_format = raw_format.strip().lower()
|
||||
if export_format not in SUPPORTED_EXPORT_FORMATS:
|
||||
supported = ", ".join(sorted(SUPPORTED_EXPORT_FORMATS))
|
||||
raise ValueError(
|
||||
f"Неподдерживаемый формат экспорта: {raw_format}. "
|
||||
f"Доступны: {supported}."
|
||||
)
|
||||
|
||||
half = _boolean(data, "half")
|
||||
int8 = _boolean(data, "int8")
|
||||
if half and int8:
|
||||
raise ValueError("Параметры half и int8 нельзя включать одновременно.")
|
||||
|
||||
return cls(
|
||||
model=_model_path(data.get("model")),
|
||||
export_format=export_format,
|
||||
imgsz=_bounded_int(data, "imgsz", 640, minimum=32, maximum=8192),
|
||||
half=half,
|
||||
int8=int8,
|
||||
dynamic=_boolean(data, "dynamic"),
|
||||
simplify=_boolean(data, "simplify"),
|
||||
batch=_bounded_int(data, "batch", 1, minimum=1, maximum=1024),
|
||||
workspace=_bounded_number(
|
||||
data,
|
||||
"workspace",
|
||||
4,
|
||||
minimum=1,
|
||||
maximum=64,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _result_path(value: Any) -> Path:
|
||||
if isinstance(value, (list, tuple)):
|
||||
if not value:
|
||||
raise RuntimeError("Ultralytics не вернул путь к экспортированной модели.")
|
||||
value = value[0]
|
||||
if not isinstance(value, (str, os.PathLike)):
|
||||
raise RuntimeError("Ultralytics не вернул путь к экспортированной модели.")
|
||||
|
||||
raw_path = os.fspath(value)
|
||||
if not raw_path or any(character in raw_path for character in ("\0", "\r", "\n")):
|
||||
raise RuntimeError("Ultralytics вернул некорректный путь результата.")
|
||||
path = Path(raw_path).expanduser()
|
||||
if not path.is_absolute():
|
||||
path = Path.cwd() / path
|
||||
try:
|
||||
return path.resolve(strict=True)
|
||||
except (FileNotFoundError, OSError, RuntimeError, ValueError) as exc:
|
||||
raise RuntimeError(
|
||||
f"Экспорт завершился без ожидаемого артефакта: {raw_path}"
|
||||
) from exc
|
||||
|
||||
|
||||
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.export_runner <config_json_path>",
|
||||
|
|
@ -16,63 +215,61 @@ def main() -> int:
|
|||
)
|
||||
return 1
|
||||
|
||||
config_path = Path(args[0])
|
||||
|
||||
# The parent waits for this marker
|
||||
# The parent waits for this marker before treating the process as started.
|
||||
print("__YOLO_WEBUI_READY__", flush=True)
|
||||
|
||||
try:
|
||||
with config_path.open("r", encoding="utf-8") as config_file:
|
||||
config_dict = json.load(config_file)
|
||||
with Path(args[0]).open("r", encoding="utf-8") as config_file:
|
||||
raw_config = json.load(config_file)
|
||||
if not isinstance(raw_config, Mapping):
|
||||
raise ValueError("Конфигурация экспорта должна быть JSON-объектом.")
|
||||
config = ExportConfig.from_mapping(raw_config)
|
||||
except Exception as exc:
|
||||
print(f"Error loading config: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
model_path = config_dict.get("model")
|
||||
export_format = config_dict.get("format", "onnx")
|
||||
imgsz = config_dict.get("imgsz", 640)
|
||||
half = config_dict.get("half", False)
|
||||
int8 = config_dict.get("int8", False)
|
||||
dynamic = config_dict.get("dynamic", False)
|
||||
simplify = config_dict.get("simplify", False)
|
||||
batch = config_dict.get("batch", 1)
|
||||
workspace = config_dict.get("workspace", 4)
|
||||
|
||||
if not model_path or not Path(model_path).exists():
|
||||
print(f"Model file not found: {model_path}")
|
||||
print(f"Ошибка конфигурации экспорта: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
try:
|
||||
os.environ["ULTRALYTICS_SAFE_LOAD"] = "1"
|
||||
from ultralytics import YOLO
|
||||
|
||||
print(f"Загрузка модели {model_path}...", flush=True)
|
||||
model = YOLO(model_path)
|
||||
print(f"Загрузка модели {config.model}...", flush=True)
|
||||
model = YOLO(str(config.model))
|
||||
|
||||
print(f"Запуск экспорта в формат {export_format}...", flush=True)
|
||||
print(f"Параметры: imgsz={imgsz}, half={half}, int8={int8}, dynamic={dynamic}, simplify={simplify}, batch={batch}, workspace={workspace}", flush=True)
|
||||
|
||||
exported_path = model.export(
|
||||
format=export_format,
|
||||
imgsz=imgsz,
|
||||
half=half,
|
||||
int8=int8,
|
||||
dynamic=dynamic,
|
||||
simplify=simplify,
|
||||
batch=batch,
|
||||
workspace=workspace
|
||||
print(
|
||||
f"Запуск экспорта в формат {config.export_format}...",
|
||||
flush=True,
|
||||
)
|
||||
print(
|
||||
"Параметры: "
|
||||
f"imgsz={config.imgsz}, half={config.half}, int8={config.int8}, "
|
||||
f"dynamic={config.dynamic}, simplify={config.simplify}, "
|
||||
f"batch={config.batch}, workspace={config.workspace}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
print(f"Экспорт завершен успешно.", flush=True)
|
||||
if exported_path:
|
||||
# ultralytics returns either a single path (string) or list of paths depending on the format.
|
||||
if isinstance(exported_path, list):
|
||||
exported_path = exported_path[0]
|
||||
print(f"__YOLO_WEBUI_RESULT__:{exported_path}", flush=True)
|
||||
export_kwargs: dict[str, Any] = {
|
||||
"format": config.export_format,
|
||||
"imgsz": config.imgsz,
|
||||
"half": config.half,
|
||||
"int8": config.int8,
|
||||
"dynamic": config.dynamic,
|
||||
"simplify": config.simplify,
|
||||
"batch": config.batch,
|
||||
}
|
||||
if config.export_format in ("engine", "tensorrt", "trt"):
|
||||
export_kwargs["workspace"] = config.workspace
|
||||
|
||||
exported_path = model.export(**export_kwargs)
|
||||
result_path = _result_path(exported_path)
|
||||
|
||||
print("Экспорт завершен успешно.", flush=True)
|
||||
print(f"__YOLO_WEBUI_RESULT__:{result_path}", flush=True)
|
||||
return 0
|
||||
except Exception as e:
|
||||
print(f"Ошибка при экспорте модели: {e}", flush=True)
|
||||
except Exception as exc:
|
||||
print(f"Ошибка при экспорте модели: {exc}", flush=True)
|
||||
traceback.print_exc()
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
|
|
|||
533
src/yolo_webui/mlflow_metrics.py
Normal file
533
src/yolo_webui/mlflow_metrics.py
Normal file
|
|
@ -0,0 +1,533 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
from collections.abc import Iterable, Mapping, MutableMapping
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .config import YoloTask
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MONITORING_SCHEMA_VERSION = "1"
|
||||
|
||||
# Public component names intentionally do not mirror Ultralytics' one-letter
|
||||
# suffixes. In particular, OBB uses "B" upstream even though its boxes are rotated.
|
||||
TASK_COMPONENTS: dict[YoloTask, tuple[tuple[str, str], ...]] = {
|
||||
"detect": (("box", "box"),),
|
||||
"segment": (("box", "box"), ("mask", "seg")),
|
||||
"classify": (),
|
||||
"pose": (("box", "box"), ("keypoints", "pose")),
|
||||
"obb": (("oriented_box", "box"),),
|
||||
}
|
||||
|
||||
PER_CLASS_FIELDS = (
|
||||
"task",
|
||||
"component",
|
||||
"class_id",
|
||||
"class_name",
|
||||
"support",
|
||||
"precision",
|
||||
"recall",
|
||||
"f1",
|
||||
"map50",
|
||||
"map50_95",
|
||||
)
|
||||
|
||||
|
||||
def _finite_float(value: Any) -> float | None:
|
||||
"""Return a finite Python float for scalar-like values."""
|
||||
if value is None or isinstance(value, (str, bytes, bool)):
|
||||
return None
|
||||
for method in ("detach", "cpu"):
|
||||
operation = getattr(value, method, None)
|
||||
if callable(operation):
|
||||
try:
|
||||
value = operation()
|
||||
except Exception:
|
||||
return None
|
||||
item = getattr(value, "item", None)
|
||||
if callable(item):
|
||||
try:
|
||||
value = item()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
number = float(value)
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return None
|
||||
return number if math.isfinite(number) else None
|
||||
|
||||
|
||||
def _finite_values(value: Any) -> list[float]:
|
||||
"""Flatten array-like values while dropping non-finite entries."""
|
||||
if value is None or isinstance(value, (str, bytes, bool)):
|
||||
return []
|
||||
for method in ("detach", "cpu"):
|
||||
operation = getattr(value, method, None)
|
||||
if callable(operation):
|
||||
try:
|
||||
value = operation()
|
||||
except Exception:
|
||||
return []
|
||||
tolist = getattr(value, "tolist", None)
|
||||
if callable(tolist):
|
||||
try:
|
||||
value = tolist()
|
||||
except Exception:
|
||||
return []
|
||||
if isinstance(value, Mapping):
|
||||
source: Iterable[Any] = value.values()
|
||||
elif isinstance(value, Iterable):
|
||||
source = value
|
||||
else:
|
||||
number = _finite_float(value)
|
||||
return [] if number is None else [number]
|
||||
|
||||
result: list[float] = []
|
||||
for item in source:
|
||||
result.extend(_finite_values(item))
|
||||
return result
|
||||
|
||||
|
||||
def _aligned_values(value: Any) -> list[float | None]:
|
||||
"""Convert a one-dimensional array without shifting non-finite positions."""
|
||||
if value is None or isinstance(value, (str, bytes, bool)):
|
||||
return []
|
||||
for method in ("detach", "cpu"):
|
||||
operation = getattr(value, method, None)
|
||||
if callable(operation):
|
||||
try:
|
||||
value = operation()
|
||||
except Exception:
|
||||
return []
|
||||
tolist = getattr(value, "tolist", None)
|
||||
if callable(tolist):
|
||||
try:
|
||||
value = tolist()
|
||||
except Exception:
|
||||
return []
|
||||
if isinstance(value, Mapping):
|
||||
source: Iterable[Any] = value.values()
|
||||
elif isinstance(value, Iterable):
|
||||
source = value
|
||||
else:
|
||||
return [_finite_float(value)]
|
||||
return [_finite_float(item) for item in source]
|
||||
|
||||
|
||||
def _attribute(obj: Any, name: str) -> Any:
|
||||
if obj is None:
|
||||
return None
|
||||
value = getattr(obj, name, None)
|
||||
if callable(value):
|
||||
try:
|
||||
return value()
|
||||
except Exception:
|
||||
return None
|
||||
return value
|
||||
|
||||
|
||||
def _add_metric(metrics: dict[str, float], key: str, value: Any) -> None:
|
||||
number = _finite_float(value)
|
||||
if number is not None:
|
||||
metrics[key] = number
|
||||
|
||||
|
||||
def _mean(values: list[float]) -> float | None:
|
||||
return sum(values) / len(values) if values else None
|
||||
|
||||
|
||||
def _harmonic_mean(precision: float | None, recall: float | None) -> float | None:
|
||||
if precision is None or recall is None:
|
||||
return None
|
||||
denominator = precision + recall
|
||||
return 0.0 if denominator == 0 else 2 * precision * recall / denominator
|
||||
|
||||
|
||||
def _component_metrics(name: str, component: Any) -> dict[str, float]:
|
||||
prefix = f"monitor/quality/{name}"
|
||||
result: dict[str, float] = {}
|
||||
precision = _finite_float(_attribute(component, "mp"))
|
||||
recall = _finite_float(_attribute(component, "mr"))
|
||||
f1_values = _finite_values(_attribute(component, "f1"))
|
||||
ap_values = _finite_values(_attribute(component, "ap"))
|
||||
mean_f1 = _mean(f1_values)
|
||||
|
||||
_add_metric(result, f"{prefix}/precision", precision)
|
||||
_add_metric(result, f"{prefix}/recall", recall)
|
||||
_add_metric(
|
||||
result,
|
||||
f"{prefix}/f1_at_optimal_confidence",
|
||||
mean_f1 if mean_f1 is not None else _harmonic_mean(precision, recall),
|
||||
)
|
||||
_add_metric(result, f"{prefix}/map50", _attribute(component, "map50"))
|
||||
_add_metric(result, f"{prefix}/map75", _attribute(component, "map75"))
|
||||
_add_metric(result, f"{prefix}/map50_95", _attribute(component, "map"))
|
||||
_add_metric(
|
||||
result,
|
||||
f"{prefix}/worst_class_f1_at_optimal_confidence",
|
||||
min(f1_values) if f1_values else None,
|
||||
)
|
||||
_add_metric(result, f"{prefix}/worst_class_map50_95", min(ap_values) if ap_values else None)
|
||||
_add_metric(result, f"{prefix}/classes_evaluated", len(ap_values) or len(f1_values))
|
||||
return result
|
||||
|
||||
|
||||
def _matrix_rows(matrix: Any) -> list[list[float]]:
|
||||
raw_rows = _attribute(matrix, "tolist")
|
||||
if raw_rows is None:
|
||||
raw_rows = matrix
|
||||
if not isinstance(raw_rows, Iterable) or isinstance(raw_rows, (str, bytes)):
|
||||
return []
|
||||
|
||||
rows: list[list[float]] = []
|
||||
for row in raw_rows:
|
||||
if not isinstance(row, Iterable) or isinstance(row, (str, bytes)):
|
||||
return []
|
||||
converted: list[float] = []
|
||||
for value in row:
|
||||
number = _finite_float(value)
|
||||
converted.append(0.0 if number is None else max(0.0, number))
|
||||
rows.append(converted)
|
||||
size = len(rows)
|
||||
return rows if size and all(len(row) == size for row in rows) else []
|
||||
|
||||
|
||||
def _classification_statistics(metric_set: Any) -> tuple[dict[str, float], list[dict[str, Any]]]:
|
||||
prefix = "monitor/quality/classification"
|
||||
result: dict[str, float] = {}
|
||||
top1 = _finite_float(_attribute(metric_set, "top1"))
|
||||
top5 = _finite_float(_attribute(metric_set, "top5"))
|
||||
_add_metric(result, f"{prefix}/top1_accuracy", top1)
|
||||
_add_metric(result, f"{prefix}/top5_accuracy", top5)
|
||||
_add_metric(result, f"{prefix}/top1_error", None if top1 is None else 1.0 - top1)
|
||||
_add_metric(result, f"{prefix}/top5_error", None if top5 is None else 1.0 - top5)
|
||||
|
||||
confusion = _attribute(metric_set, "confusion_matrix")
|
||||
rows = _matrix_rows(_attribute(confusion, "matrix"))
|
||||
names = _attribute(confusion, "names") or {}
|
||||
per_class: list[dict[str, Any]] = []
|
||||
if not rows:
|
||||
return result, per_class
|
||||
|
||||
precisions: list[float] = []
|
||||
supported_recalls: list[float] = []
|
||||
f1_scores: list[float] = []
|
||||
supports: list[float] = []
|
||||
for class_id in range(len(rows)):
|
||||
true_positive = rows[class_id][class_id]
|
||||
predicted = sum(rows[class_id])
|
||||
support = sum(row[class_id] for row in rows)
|
||||
if support <= 0 and predicted <= 0:
|
||||
continue
|
||||
precision = true_positive / predicted if predicted else 0.0
|
||||
recall = true_positive / support if support else 0.0
|
||||
f1 = _harmonic_mean(precision, recall) or 0.0
|
||||
precisions.append(precision)
|
||||
f1_scores.append(f1)
|
||||
supports.append(support)
|
||||
if support > 0:
|
||||
supported_recalls.append(recall)
|
||||
per_class.append(
|
||||
{
|
||||
"task": "classify",
|
||||
"component": "classification",
|
||||
"class_id": class_id,
|
||||
"class_name": _class_name(names, class_id),
|
||||
"support": support,
|
||||
"precision": precision,
|
||||
"recall": recall,
|
||||
"f1": f1,
|
||||
"map50": "",
|
||||
"map50_95": "",
|
||||
}
|
||||
)
|
||||
|
||||
total_support = sum(supports)
|
||||
weighted_f1 = (
|
||||
sum(score * support for score, support in zip(f1_scores, supports)) / total_support
|
||||
if total_support
|
||||
else None
|
||||
)
|
||||
_add_metric(result, f"{prefix}/macro_precision", _mean(precisions))
|
||||
_add_metric(result, f"{prefix}/macro_recall", _mean(supported_recalls))
|
||||
_add_metric(result, f"{prefix}/macro_f1", _mean(f1_scores))
|
||||
_add_metric(result, f"{prefix}/weighted_f1", weighted_f1)
|
||||
_add_metric(result, f"{prefix}/balanced_accuracy", _mean(supported_recalls))
|
||||
_add_metric(
|
||||
result,
|
||||
f"{prefix}/worst_class_recall",
|
||||
min(supported_recalls) if supported_recalls else None,
|
||||
)
|
||||
_add_metric(result, f"{prefix}/classes_evaluated", len(per_class))
|
||||
return result, per_class
|
||||
|
||||
|
||||
def collect_monitoring_metrics(task: YoloTask, trainer: Any) -> dict[str, float]:
|
||||
"""Collect stable, task-aware metrics from an Ultralytics trainer."""
|
||||
result: dict[str, float] = {}
|
||||
trainer_metrics = getattr(trainer, "metrics", {}) or {}
|
||||
validator = getattr(trainer, "validator", None)
|
||||
metric_set = getattr(validator, "metrics", None)
|
||||
final_validation = (
|
||||
validator is not None and getattr(validator, "training", True) is False
|
||||
)
|
||||
|
||||
current_fitness = _finite_float(_attribute(metric_set, "fitness"))
|
||||
if current_fitness is None:
|
||||
current_fitness = _finite_float(getattr(trainer, "fitness", None))
|
||||
best_fitness = _finite_float(getattr(trainer, "best_fitness", None))
|
||||
_add_metric(result, "monitor/fitness/ultralytics_current", current_fitness)
|
||||
_add_metric(result, "monitor/fitness/ultralytics_best", best_fitness)
|
||||
fitness_scale = max(1, len(TASK_COMPONENTS[task]))
|
||||
_add_metric(
|
||||
result,
|
||||
"monitor/fitness/task_score_current",
|
||||
None if current_fitness is None else current_fitness / fitness_scale,
|
||||
)
|
||||
_add_metric(
|
||||
result,
|
||||
"monitor/fitness/task_score_best",
|
||||
None if best_fitness is None else best_fitness / fitness_scale,
|
||||
)
|
||||
|
||||
if not final_validation:
|
||||
train_total: float | None = None
|
||||
validation_total: float | None = None
|
||||
label_losses = getattr(trainer, "label_loss_items", None)
|
||||
if callable(label_losses):
|
||||
try:
|
||||
train_losses = label_losses(getattr(trainer, "tloss", None), prefix="train")
|
||||
except Exception:
|
||||
train_losses = {}
|
||||
if isinstance(train_losses, Mapping):
|
||||
values = [
|
||||
number
|
||||
for value in train_losses.values()
|
||||
if (number := _finite_float(value)) is not None
|
||||
]
|
||||
train_total = sum(values) if values else None
|
||||
_add_metric(result, "monitor/loss/train_total", train_total)
|
||||
|
||||
if isinstance(trainer_metrics, Mapping):
|
||||
validation_loss_keys: set[str] = set()
|
||||
if callable(label_losses):
|
||||
try:
|
||||
expected_losses = label_losses(None, prefix="val")
|
||||
except Exception:
|
||||
expected_losses = ()
|
||||
if isinstance(expected_losses, Mapping):
|
||||
validation_loss_keys.update(map(str, expected_losses))
|
||||
elif isinstance(expected_losses, Iterable) and not isinstance(
|
||||
expected_losses,
|
||||
(str, bytes),
|
||||
):
|
||||
validation_loss_keys.update(map(str, expected_losses))
|
||||
if not validation_loss_keys:
|
||||
validation_loss_keys = {
|
||||
str(key)
|
||||
for key in trainer_metrics
|
||||
if str(key) == "val/loss"
|
||||
or (
|
||||
str(key).startswith("val/")
|
||||
and str(key).endswith("_loss")
|
||||
)
|
||||
}
|
||||
val_losses = [
|
||||
number
|
||||
for key, value in trainer_metrics.items()
|
||||
if str(key) in validation_loss_keys
|
||||
and (number := _finite_float(value)) is not None
|
||||
]
|
||||
validation_total = sum(val_losses) if val_losses else None
|
||||
_add_metric(result, "monitor/loss/validation_total", validation_total)
|
||||
|
||||
if train_total is not None and validation_total is not None:
|
||||
_add_metric(
|
||||
result,
|
||||
"monitor/loss/generalization_gap",
|
||||
validation_total - train_total,
|
||||
)
|
||||
if train_total > 0:
|
||||
_add_metric(
|
||||
result,
|
||||
"monitor/loss/validation_to_train_ratio",
|
||||
validation_total / train_total,
|
||||
)
|
||||
|
||||
learning_rates = _finite_values(getattr(trainer, "lr", None))
|
||||
_add_metric(result, "monitor/optimization/learning_rate_mean", _mean(learning_rates))
|
||||
_add_metric(
|
||||
result,
|
||||
"monitor/performance/epoch_seconds",
|
||||
getattr(trainer, "epoch_time", None),
|
||||
)
|
||||
|
||||
speed = _attribute(metric_set, "speed")
|
||||
if isinstance(speed, Mapping):
|
||||
for stage in ("preprocess", "inference", "loss", "postprocess"):
|
||||
_add_metric(
|
||||
result,
|
||||
f"monitor/performance/validation_{stage}_ms_per_image",
|
||||
speed.get(stage),
|
||||
)
|
||||
|
||||
if task == "classify":
|
||||
classification, _ = _classification_statistics(metric_set)
|
||||
result.update(classification)
|
||||
else:
|
||||
for public_name, attribute_name in TASK_COMPONENTS[task]:
|
||||
component = _attribute(metric_set, attribute_name)
|
||||
if component is not None:
|
||||
result.update(_component_metrics(public_name, component))
|
||||
return result
|
||||
|
||||
|
||||
def _class_name(names: Any, class_id: int) -> str:
|
||||
if isinstance(names, Mapping):
|
||||
return str(names.get(class_id, names.get(str(class_id), class_id)))
|
||||
if isinstance(names, (list, tuple)) and 0 <= class_id < len(names):
|
||||
return str(names[class_id])
|
||||
return str(class_id)
|
||||
|
||||
|
||||
def _value_at(values: list[float | None], index: int) -> float | str:
|
||||
if index >= len(values) or values[index] is None:
|
||||
return ""
|
||||
return values[index]
|
||||
|
||||
|
||||
def _component_rows(
|
||||
task: YoloTask,
|
||||
public_name: str,
|
||||
component: Any,
|
||||
metric_set: Any,
|
||||
) -> list[dict[str, Any]]:
|
||||
raw_class_indices = _aligned_values(_attribute(component, "ap_class_index"))
|
||||
precision = _aligned_values(_attribute(component, "p"))
|
||||
recall = _aligned_values(_attribute(component, "r"))
|
||||
f1 = _aligned_values(_attribute(component, "f1"))
|
||||
map50 = _aligned_values(_attribute(component, "ap50"))
|
||||
map50_95 = _aligned_values(_attribute(component, "ap"))
|
||||
count = max(
|
||||
map(len, (raw_class_indices, precision, recall, f1, map50, map50_95)),
|
||||
default=0,
|
||||
)
|
||||
class_indices = [
|
||||
index if value is None else int(value)
|
||||
for index, value in enumerate(raw_class_indices)
|
||||
]
|
||||
if not class_indices:
|
||||
class_indices = list(range(count))
|
||||
|
||||
names = _attribute(metric_set, "names") or {}
|
||||
support_by_class = _aligned_values(_attribute(metric_set, "nt_per_class"))
|
||||
rows: list[dict[str, Any]] = []
|
||||
for index in range(min(count, len(class_indices))):
|
||||
class_id = class_indices[index]
|
||||
rows.append(
|
||||
{
|
||||
"task": task,
|
||||
"component": public_name,
|
||||
"class_id": class_id,
|
||||
"class_name": _class_name(names, class_id),
|
||||
"support": _value_at(support_by_class, class_id),
|
||||
"precision": _value_at(precision, index),
|
||||
"recall": _value_at(recall, index),
|
||||
"f1": _value_at(f1, index),
|
||||
"map50": _value_at(map50, index),
|
||||
"map50_95": _value_at(map50_95, index),
|
||||
}
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def collect_per_class_metrics(task: YoloTask, trainer: Any) -> list[dict[str, Any]]:
|
||||
"""Build final per-class diagnostics for the MLflow CSV artifact."""
|
||||
validator = getattr(trainer, "validator", None)
|
||||
metric_set = getattr(validator, "metrics", None)
|
||||
if metric_set is None:
|
||||
return []
|
||||
if task == "classify":
|
||||
_, rows = _classification_statistics(metric_set)
|
||||
return rows
|
||||
|
||||
result: list[dict[str, Any]] = []
|
||||
for public_name, attribute_name in TASK_COMPONENTS[task]:
|
||||
component = _attribute(metric_set, attribute_name)
|
||||
if component is not None:
|
||||
result.extend(_component_rows(task, public_name, component, metric_set))
|
||||
return result
|
||||
|
||||
|
||||
class TaskMetricsMonitor:
|
||||
"""Enrich Ultralytics metrics before its built-in MLflow callback runs."""
|
||||
|
||||
def __init__(self, task: YoloTask, *, mlflow_enabled: bool) -> None:
|
||||
self.task = task
|
||||
self.mlflow_enabled = mlflow_enabled
|
||||
|
||||
def on_train_start(self, trainer: Any) -> None:
|
||||
"""Attach searchable task/schema tags without ever starting a second run."""
|
||||
if not self.mlflow_enabled or not getattr(trainer, "_mlflow_active", False):
|
||||
return
|
||||
try:
|
||||
import mlflow
|
||||
|
||||
if mlflow.active_run() is not None:
|
||||
tags = {
|
||||
"monitoring.schema_version": MONITORING_SCHEMA_VERSION,
|
||||
"yolo.task": self.task,
|
||||
}
|
||||
run_group = os.environ.get("YOLO_WEBUI_MLFLOW_RUN_GROUP", "").strip()
|
||||
if run_group:
|
||||
tags["yolo.run_group"] = run_group
|
||||
mlflow.set_tags(tags)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Не удалось записать теги мониторинга "
|
||||
"в MLflow: %s",
|
||||
exc,
|
||||
)
|
||||
|
||||
def on_fit_epoch_end(self, trainer: Any) -> None:
|
||||
"""Add derived metrics to trainer.metrics for the current validation epoch."""
|
||||
if not self.mlflow_enabled:
|
||||
return
|
||||
metrics = getattr(trainer, "metrics", None)
|
||||
if not isinstance(metrics, MutableMapping):
|
||||
return
|
||||
try:
|
||||
metrics.update(collect_monitoring_metrics(self.task, trainer))
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Не удалось собрать task-specific метрики: %s",
|
||||
exc,
|
||||
)
|
||||
|
||||
def on_train_end(self, trainer: Any) -> None:
|
||||
"""Write and upload final per-class diagnostics to the active MLflow run."""
|
||||
if not self.mlflow_enabled:
|
||||
return
|
||||
try:
|
||||
rows = collect_per_class_metrics(self.task, trainer)
|
||||
output = Path(trainer.save_dir) / "monitoring" / "task_metrics.csv"
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
with output.open("w", newline="", encoding="utf-8") as stream:
|
||||
writer = csv.DictWriter(stream, fieldnames=PER_CLASS_FIELDS)
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
if getattr(trainer, "_mlflow_active", False):
|
||||
import mlflow
|
||||
|
||||
if mlflow.active_run() is not None:
|
||||
mlflow.log_artifact(str(output), artifact_path="monitoring")
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Не удалось сохранить per-class метрики: %s",
|
||||
exc,
|
||||
)
|
||||
|
|
@ -24,6 +24,7 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
const mlflowInputs = document.querySelectorAll('.mlflow-fields input');
|
||||
const trackingUriInput = document.getElementById('tracking-uri');
|
||||
const mlflowHeaderLink = document.getElementById('mlflow-header-link');
|
||||
const taskSelect = document.getElementById('task');
|
||||
|
||||
// --- Dynamic Model Selection ---
|
||||
const standardModels = {
|
||||
|
|
@ -34,6 +35,11 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
obb: ['yolo11n-obb.pt', 'yolo11s-obb.pt', 'yolo11m-obb.pt', 'yolo11l-obb.pt', 'yolo11x-obb.pt']
|
||||
};
|
||||
let discoveredModels = [];
|
||||
let discoveredDatasets = [];
|
||||
|
||||
function getDiscoveredModelValue(model) {
|
||||
return model.path || model.name;
|
||||
}
|
||||
|
||||
function updateModelOptions() {
|
||||
const task = taskSelect.value;
|
||||
|
|
@ -54,13 +60,13 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
modelSelect.appendChild(stdGroup);
|
||||
|
||||
// Group 2: Discovered Models
|
||||
const localModels = discoveredModels.filter(m => !stdModels.includes(m.name));
|
||||
const localModels = discoveredModels.filter(m => getDiscoveredModelValue(m));
|
||||
if (localModels.length > 0) {
|
||||
const localGroup = document.createElement('optgroup');
|
||||
localGroup.label = 'Локальные/скачанные модели';
|
||||
localModels.forEach(m => {
|
||||
const opt = document.createElement('option');
|
||||
opt.value = m.name;
|
||||
opt.value = getDiscoveredModelValue(m);
|
||||
opt.textContent = m.name;
|
||||
localGroup.appendChild(opt);
|
||||
});
|
||||
|
|
@ -74,8 +80,8 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
modelSelect.appendChild(customOpt);
|
||||
|
||||
// Match selection if valid
|
||||
const allAvailable = [...stdModels, ...localModels.map(m => m.name)];
|
||||
if (allAvailable.includes(currentSelectVal)) {
|
||||
const allAvailable = [...stdModels, ...localModels.map(getDiscoveredModelValue)];
|
||||
if (currentSelectVal === '__custom__' || allAvailable.includes(currentSelectVal)) {
|
||||
modelSelect.value = currentSelectVal;
|
||||
} else {
|
||||
modelSelect.value = stdModels[0] || '__custom__';
|
||||
|
|
@ -95,7 +101,6 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
}
|
||||
}
|
||||
|
||||
const taskSelect = document.getElementById('task');
|
||||
|
||||
// Session Controls
|
||||
const sessionSelect = document.getElementById('session-select');
|
||||
|
|
@ -121,15 +126,15 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
const autoscrollCheck = document.getElementById('autoscroll');
|
||||
const clearLogBtn = document.getElementById('clear-log-btn');
|
||||
|
||||
// Chart
|
||||
const ctx = document.getElementById('metricsChart').getContext('2d');
|
||||
let metricsChart = null;
|
||||
|
||||
// State variables
|
||||
let trainingTimer = null;
|
||||
let secondsElapsed = 0;
|
||||
let socket = null;
|
||||
let socketReconnectTimer = null;
|
||||
let isTrainingActive = false;
|
||||
let isTrainingStartPending = false;
|
||||
let trainingStatusRevision = 0;
|
||||
let currentTrainingStatus = 'idle';
|
||||
|
||||
// --- View Tab Switching ---
|
||||
const viewBtns = document.querySelectorAll('.view-btn');
|
||||
|
|
@ -215,7 +220,17 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
function updateMlflowHeaderLink() {
|
||||
const uri = trackingUriInput.value.trim();
|
||||
if (uri.startsWith('http://') || uri.startsWith('https://')) {
|
||||
mlflowHeaderLink.href = uri;
|
||||
let browserUri = uri;
|
||||
try {
|
||||
const parsedUri = new URL(uri);
|
||||
if (parsedUri.hostname === 'mlflow') {
|
||||
parsedUri.hostname = window.location.hostname || 'localhost';
|
||||
browserUri = parsedUri.href;
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('Invalid MLflow tracking URI:', error);
|
||||
}
|
||||
mlflowHeaderLink.href = browserUri;
|
||||
mlflowHeaderLink.style.opacity = '1';
|
||||
mlflowHeaderLink.style.pointerEvents = 'auto';
|
||||
} else {
|
||||
|
|
@ -235,6 +250,11 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
mlflowEnabled.addEventListener('change', updateMlflowFields);
|
||||
|
||||
// --- Logger ---
|
||||
function scrollToLogBottom(container) {
|
||||
if (!container) return;
|
||||
container.scrollTop = container.scrollHeight;
|
||||
}
|
||||
|
||||
function addLogLine(message, level = 'info') {
|
||||
const line = document.createElement('div');
|
||||
line.className = `log-line log-level-${level.toLowerCase()}`;
|
||||
|
|
@ -242,7 +262,7 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
logContainer.appendChild(line);
|
||||
|
||||
if (autoscrollCheck.checked) {
|
||||
logContainer.scrollTop = logContainer.scrollHeight;
|
||||
scrollToLogBottom(logContainer);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -250,78 +270,6 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
logContainer.innerHTML = '';
|
||||
});
|
||||
|
||||
// --- Chart.js Integration ---
|
||||
function initChart(datasets = []) {
|
||||
if (metricsChart) {
|
||||
metricsChart.destroy();
|
||||
}
|
||||
|
||||
metricsChart = new Chart(ctx, {
|
||||
type: 'line',
|
||||
data: {
|
||||
labels: [],
|
||||
datasets: datasets
|
||||
},
|
||||
options: {
|
||||
responsive: true,
|
||||
maintainAspectRatio: false,
|
||||
scales: {
|
||||
x: {
|
||||
title: { display: true, text: 'Эпоха', color: '#a1a1aa' },
|
||||
grid: { color: '#27272a' },
|
||||
ticks: { color: '#a1a1aa' }
|
||||
},
|
||||
y: {
|
||||
title: { display: true, text: 'Значение', color: '#a1a1aa' },
|
||||
grid: { color: '#27272a' },
|
||||
ticks: { color: '#a1a1aa' }
|
||||
}
|
||||
},
|
||||
plugins: {
|
||||
legend: {
|
||||
labels: { color: '#f4f4f5', font: { family: 'Outfit' } }
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
function updateChart(epoch, metrics) {
|
||||
if (!metricsChart) {
|
||||
initChart();
|
||||
}
|
||||
|
||||
let labelIndex = metricsChart.data.labels.indexOf(epoch);
|
||||
if (labelIndex === -1) {
|
||||
metricsChart.data.labels.push(epoch);
|
||||
labelIndex = metricsChart.data.labels.length - 1;
|
||||
metricsChart.data.datasets.forEach(dataset => dataset.data.push(null));
|
||||
}
|
||||
|
||||
const colors = ['#f97316', '#10b981', '#3b82f6', '#eab308', '#a855f7'];
|
||||
Object.entries(metrics).forEach(([key, value]) => {
|
||||
if (key === 'epoch') return;
|
||||
|
||||
let dataset = metricsChart.data.datasets.find(item => item.label === key);
|
||||
if (!dataset) {
|
||||
const color = colors[metricsChart.data.datasets.length % colors.length];
|
||||
dataset = {
|
||||
label: key,
|
||||
data: Array(metricsChart.data.labels.length).fill(null),
|
||||
borderColor: color,
|
||||
backgroundColor: color + '22',
|
||||
tension: 0.15,
|
||||
fill: false
|
||||
};
|
||||
metricsChart.data.datasets.push(dataset);
|
||||
}
|
||||
|
||||
dataset.data[labelIndex] = value;
|
||||
});
|
||||
|
||||
metricsChart.update();
|
||||
}
|
||||
|
||||
// --- Timer UI ---
|
||||
function startTimer() {
|
||||
stopTimer();
|
||||
|
|
@ -343,46 +291,72 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
}
|
||||
|
||||
// --- WebSocket Sync ---
|
||||
function scheduleWebSocketReconnect() {
|
||||
if (socketReconnectTimer !== null) return;
|
||||
socketReconnectTimer = setTimeout(() => {
|
||||
socketReconnectTimer = null;
|
||||
connectWebSocket();
|
||||
}, 5000);
|
||||
}
|
||||
|
||||
function parseWebSocketMessage(event, label, logLine) {
|
||||
try {
|
||||
const data = JSON.parse(event.data);
|
||||
if (!data || typeof data !== 'object' || Array.isArray(data)) {
|
||||
throw new Error('WebSocket payload must be an object');
|
||||
}
|
||||
return data;
|
||||
} catch (error) {
|
||||
console.error(`${label} message error:`, error);
|
||||
logLine('Получено некорректное сообщение от сервера.', 'warning');
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
function connectWebSocket() {
|
||||
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
|
||||
const wsUrl = `${protocol}//${window.location.host}/api/ws`;
|
||||
|
||||
socket = new WebSocket(wsUrl);
|
||||
let nextSocket;
|
||||
try {
|
||||
nextSocket = new WebSocket(wsUrl);
|
||||
} catch (error) {
|
||||
console.error('WS connection error:', error);
|
||||
addLogLine('Не удалось подключиться к серверу. Повторная попытка через 5 секунд...', 'warning');
|
||||
scheduleWebSocketReconnect();
|
||||
return;
|
||||
}
|
||||
socket = nextSocket;
|
||||
|
||||
socket.onopen = () => {
|
||||
nextSocket.onopen = () => {
|
||||
addLogLine('Соединение с сервером установлено.', 'info');
|
||||
};
|
||||
|
||||
socket.onclose = () => {
|
||||
nextSocket.onclose = () => {
|
||||
if (socket !== nextSocket) return;
|
||||
socket = null;
|
||||
addLogLine('Соединение потеряно. Повторная попытка через 5 секунд...', 'warning');
|
||||
setTimeout(connectWebSocket, 5000);
|
||||
scheduleWebSocketReconnect();
|
||||
};
|
||||
|
||||
socket.onerror = (err) => {
|
||||
nextSocket.onerror = (err) => {
|
||||
console.error('WS Error:', err);
|
||||
};
|
||||
|
||||
socket.onmessage = (event) => {
|
||||
const data = JSON.parse(event.data);
|
||||
nextSocket.onmessage = (event) => {
|
||||
const data = parseWebSocketMessage(event, 'Training WebSocket', addLogLine);
|
||||
if (!data) return;
|
||||
|
||||
if (data.type === 'init') {
|
||||
updateUIStatus(data.status);
|
||||
|
||||
// Load logs
|
||||
logContainer.innerHTML = '';
|
||||
data.logs.forEach(([levelCode, msg]) => {
|
||||
(Array.isArray(data.logs) ? data.logs : []).forEach(([levelCode, msg]) => {
|
||||
const level = levelCode.replace('__LOG_LEVEL_', '').replace('__', '').toLowerCase();
|
||||
addLogLine(msg, level);
|
||||
});
|
||||
|
||||
// Draw initial chart points
|
||||
initChart();
|
||||
if (data.metrics && data.metrics.length > 0) {
|
||||
data.metrics.forEach(m => {
|
||||
updateChart(m.epoch, m);
|
||||
});
|
||||
}
|
||||
|
||||
// Sync progress
|
||||
if (data.status === 'training' || data.status === 'stopping') {
|
||||
updateProgress(data.epoch, data.total_epochs);
|
||||
|
|
@ -412,18 +386,17 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
addLogLine(data.message, data.level);
|
||||
}
|
||||
if (autoscrollCheck && autoscrollCheck.checked) {
|
||||
logContainer.scrollTop = logContainer.scrollHeight;
|
||||
scrollToLogBottom(logContainer);
|
||||
}
|
||||
} else if (data.type === 'progress') {
|
||||
updateProgress(data.epoch, data.total_epochs, data.message);
|
||||
if (data.metrics) {
|
||||
updateChart(data.epoch, data.metrics);
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
function updateUIStatus(status) {
|
||||
trainingStatusRevision++;
|
||||
currentTrainingStatus = status;
|
||||
statusCard.className = `status-${status}`;
|
||||
|
||||
switch (status) {
|
||||
|
|
@ -442,7 +415,6 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
startBtn.disabled = true;
|
||||
stopBtn.disabled = false;
|
||||
startTimer();
|
||||
initChart();
|
||||
break;
|
||||
case 'training':
|
||||
statusTitle.textContent = 'ОБУЧЕНИЕ';
|
||||
|
|
@ -505,6 +477,25 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
return Number.isNaN(value) ? fallback : value;
|
||||
}
|
||||
|
||||
function readValidatedNumber(id, label, {integer = false, min, max} = {}) {
|
||||
const rawValue = document.getElementById(id).value.trim();
|
||||
const value = Number(rawValue);
|
||||
const outOfRange = (min !== undefined && value < min)
|
||||
|| (max !== undefined && value > max);
|
||||
if (
|
||||
rawValue === ''
|
||||
|| !Number.isFinite(value)
|
||||
|| (integer && !Number.isInteger(value))
|
||||
|| outOfRange
|
||||
) {
|
||||
const range = max === undefined ? `не меньше ${min}` : `от ${min} до ${max}`;
|
||||
const integerHint = integer ? 'целым числом ' : '';
|
||||
showNotification(`${label} должен быть ${integerHint}${range}.`, 'warning');
|
||||
return null;
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
function getFormConfig() {
|
||||
return {
|
||||
dataset: document.getElementById('dataset').value.trim(),
|
||||
|
|
@ -561,11 +552,18 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
const modelVal = data.model || 'yolo11n.pt';
|
||||
const task = data.task || 'detect';
|
||||
const stdModels = standardModels[task] || [];
|
||||
const allAvailable = [...stdModels, ...discoveredModels.map(m => m.name)];
|
||||
if (allAvailable.includes(modelVal)) {
|
||||
const matchedLocalModel = discoveredModels.find(
|
||||
model => model.path === modelVal || model.name === modelVal
|
||||
);
|
||||
if (stdModels.includes(modelVal)) {
|
||||
modelSelect.value = modelVal;
|
||||
modelCustomWrapper.style.display = 'none';
|
||||
document.getElementById('model').value = modelVal;
|
||||
} else if (matchedLocalModel) {
|
||||
const discoveredValue = getDiscoveredModelValue(matchedLocalModel);
|
||||
modelSelect.value = discoveredValue;
|
||||
modelCustomWrapper.style.display = 'none';
|
||||
document.getElementById('model').value = discoveredValue;
|
||||
} else {
|
||||
modelSelect.value = '__custom__';
|
||||
modelCustomWrapper.style.display = 'block';
|
||||
|
|
@ -709,7 +707,7 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
}
|
||||
|
||||
// --- Datasets Auto-Discovery ---
|
||||
let discoveredDatasets = [];
|
||||
|
||||
|
||||
async function loadDatasetsList() {
|
||||
try {
|
||||
|
|
@ -831,9 +829,9 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
});
|
||||
|
||||
sessionSaveBtn.addEventListener('click', async () => {
|
||||
const name = sessionNameInput.value.trim().replace(/[^a-zA-Z0-9_\-]/g, "");
|
||||
if (!name) {
|
||||
showNotification('Введите корректное имя профиля (латиница, цифры, дефисы).', 'warning');
|
||||
const name = sessionNameInput.value.trim();
|
||||
if (!/^[a-zA-Z0-9_-]+$/.test(name)) {
|
||||
showNotification('Введите корректное имя профиля (латиница, цифры, дефисы и подчёркивания).', 'warning');
|
||||
return;
|
||||
}
|
||||
if (name === "last_run") {
|
||||
|
|
@ -888,8 +886,12 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
|
||||
// --- Form submit ---
|
||||
async function startTraining() {
|
||||
if (isTrainingActive) return;
|
||||
if (isTrainingActive || isTrainingStartPending) return;
|
||||
isTrainingStartPending = true;
|
||||
startBtn.disabled = true;
|
||||
const config = getFormConfig();
|
||||
const statusRevisionAtStart = trainingStatusRevision;
|
||||
let started = false;
|
||||
|
||||
try {
|
||||
const res = await fetch('/api/train/start', {
|
||||
|
|
@ -902,10 +904,21 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
if (!res.ok) {
|
||||
throw new Error(data.detail || 'Failed to start training');
|
||||
}
|
||||
started = true;
|
||||
const terminalStatusReceived = trainingStatusRevision !== statusRevisionAtStart
|
||||
&& ['finished', 'succeeded', 'cancelled', 'failed'].includes(currentTrainingStatus);
|
||||
if (!isTrainingActive && !terminalStatusReceived) {
|
||||
updateUIStatus('preparing');
|
||||
}
|
||||
showNotification('Обучение успешно запущено!', 'success');
|
||||
} catch (err) {
|
||||
console.error('Start error:', err);
|
||||
showNotification(err.message, 'error');
|
||||
} finally {
|
||||
isTrainingStartPending = false;
|
||||
if (!started && !isTrainingActive) {
|
||||
startBtn.disabled = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -925,6 +938,7 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
}
|
||||
|
||||
if (configForm) {
|
||||
configForm.addEventListener('submit', (e) => e.preventDefault());
|
||||
configForm.addEventListener('input', () => {
|
||||
const config = getFormConfig();
|
||||
localStorage.setItem('draft_config', JSON.stringify(config));
|
||||
|
|
@ -1002,7 +1016,11 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
const exportAutoscrollCheck = document.getElementById('export-autoscroll');
|
||||
|
||||
let isExportActive = false;
|
||||
let isExportStartPending = false;
|
||||
let exportSocket = null;
|
||||
let exportSocketReconnectTimer = null;
|
||||
let exportStatusRevision = 0;
|
||||
let currentExportStatus = 'idle';
|
||||
|
||||
function addExportLogLine(message, level = 'info') {
|
||||
const line = document.createElement('div');
|
||||
|
|
@ -1011,7 +1029,7 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
exportLogContainer.appendChild(line);
|
||||
|
||||
if (exportAutoscrollCheck.checked) {
|
||||
exportLogContainer.scrollTop = exportLogContainer.scrollHeight;
|
||||
scrollToLogBottom(exportLogContainer);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1019,32 +1037,52 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
exportLogContainer.innerHTML = '';
|
||||
});
|
||||
|
||||
function scheduleExportWebSocketReconnect() {
|
||||
if (exportSocketReconnectTimer !== null) return;
|
||||
exportSocketReconnectTimer = setTimeout(() => {
|
||||
exportSocketReconnectTimer = null;
|
||||
connectExportWebSocket();
|
||||
}, 5000);
|
||||
}
|
||||
|
||||
function connectExportWebSocket() {
|
||||
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
|
||||
const wsUrl = `${protocol}//${window.location.host}/api/export/ws`;
|
||||
|
||||
exportSocket = new WebSocket(wsUrl);
|
||||
let nextSocket;
|
||||
try {
|
||||
nextSocket = new WebSocket(wsUrl);
|
||||
} catch (error) {
|
||||
console.error('Export WS connection error:', error);
|
||||
addExportLogLine('Не удалось подключиться к серверу. Повторная попытка через 5 секунд...', 'warning');
|
||||
scheduleExportWebSocketReconnect();
|
||||
return;
|
||||
}
|
||||
exportSocket = nextSocket;
|
||||
|
||||
exportSocket.onopen = () => {
|
||||
nextSocket.onopen = () => {
|
||||
addExportLogLine('Соединение с сервером установлено.', 'info');
|
||||
};
|
||||
|
||||
exportSocket.onclose = () => {
|
||||
nextSocket.onclose = () => {
|
||||
if (exportSocket !== nextSocket) return;
|
||||
exportSocket = null;
|
||||
addExportLogLine('Соединение потеряно. Повторная попытка через 5 секунд...', 'warning');
|
||||
setTimeout(connectExportWebSocket, 5000);
|
||||
scheduleExportWebSocketReconnect();
|
||||
};
|
||||
|
||||
exportSocket.onerror = (err) => {
|
||||
nextSocket.onerror = (err) => {
|
||||
console.error('Export WS Error:', err);
|
||||
};
|
||||
|
||||
exportSocket.onmessage = (event) => {
|
||||
const data = JSON.parse(event.data);
|
||||
nextSocket.onmessage = (event) => {
|
||||
const data = parseWebSocketMessage(event, 'Export WebSocket', addExportLogLine);
|
||||
if (!data) return;
|
||||
|
||||
if (data.type === 'init') {
|
||||
updateExportUIStatus(data.status);
|
||||
exportLogContainer.innerHTML = '';
|
||||
data.logs.forEach(([levelCode, msg]) => {
|
||||
(Array.isArray(data.logs) ? data.logs : []).forEach(([levelCode, msg]) => {
|
||||
const level = levelCode.replace('__LOG_LEVEL_', '').replace('__', '').toLowerCase();
|
||||
addExportLogLine(msg, level);
|
||||
});
|
||||
|
|
@ -1060,6 +1098,8 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
}
|
||||
|
||||
function updateExportUIStatus(status) {
|
||||
exportStatusRevision++;
|
||||
currentExportStatus = status;
|
||||
exportStatusCard.className = `status-${status}`;
|
||||
|
||||
switch (status) {
|
||||
|
|
@ -1110,7 +1150,7 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
}
|
||||
|
||||
async function startExport() {
|
||||
if (isExportActive) return;
|
||||
if (isExportActive || isExportStartPending) return;
|
||||
|
||||
const modelVal = document.getElementById('export-model').value.trim();
|
||||
if (!modelVal) {
|
||||
|
|
@ -1118,17 +1158,47 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
return;
|
||||
}
|
||||
|
||||
const imgsz = readValidatedNumber(
|
||||
'export-imgsz',
|
||||
'Размер изображения',
|
||||
{integer: true, min: 32, max: 8192}
|
||||
);
|
||||
if (imgsz === null) return;
|
||||
const batch = readValidatedNumber(
|
||||
'export-batch',
|
||||
'Размер батча',
|
||||
{integer: true, min: 1, max: 1024}
|
||||
);
|
||||
if (batch === null) return;
|
||||
const workspace = readValidatedNumber(
|
||||
'export-workspace',
|
||||
'Workspace',
|
||||
{min: 1, max: 64}
|
||||
);
|
||||
if (workspace === null) return;
|
||||
|
||||
const half = document.getElementById('export-half').checked;
|
||||
const int8 = document.getElementById('export-int8').checked;
|
||||
if (half && int8) {
|
||||
showNotification('FP16 и INT8 нельзя включать одновременно.', 'warning');
|
||||
return;
|
||||
}
|
||||
|
||||
isExportStartPending = true;
|
||||
exportStartBtn.disabled = true;
|
||||
const config = {
|
||||
model: modelVal,
|
||||
format: document.getElementById('export-format').value,
|
||||
imgsz: readNumber('export-imgsz', 640, true),
|
||||
half: document.getElementById('export-half').checked,
|
||||
int8: document.getElementById('export-int8').checked,
|
||||
imgsz,
|
||||
half,
|
||||
int8,
|
||||
dynamic: document.getElementById('export-dynamic').checked,
|
||||
simplify: document.getElementById('export-simplify').checked,
|
||||
batch: readNumber('export-batch', 1, true),
|
||||
workspace: readNumber('export-workspace', 4, true)
|
||||
batch,
|
||||
workspace
|
||||
};
|
||||
const statusRevisionAtStart = exportStatusRevision;
|
||||
let started = false;
|
||||
|
||||
try {
|
||||
const res = await fetch('/api/export/start', {
|
||||
|
|
@ -1141,10 +1211,21 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
if (!res.ok) {
|
||||
throw new Error(data.detail || 'Failed to start export');
|
||||
}
|
||||
started = true;
|
||||
const terminalStatusReceived = exportStatusRevision !== statusRevisionAtStart
|
||||
&& ['succeeded', 'cancelled', 'failed'].includes(currentExportStatus);
|
||||
if (!isExportActive && !terminalStatusReceived) {
|
||||
updateExportUIStatus('preparing');
|
||||
}
|
||||
showNotification('Экспорт успешно запущен!', 'success');
|
||||
} catch (err) {
|
||||
console.error('Export start error:', err);
|
||||
showNotification(err.message, 'error');
|
||||
} finally {
|
||||
isExportStartPending = false;
|
||||
if (!started && !isExportActive) {
|
||||
exportStartBtn.disabled = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1175,6 +1256,8 @@ document.addEventListener('DOMContentLoaded', () => {
|
|||
loadSessionsList();
|
||||
connectWebSocket();
|
||||
connectExportWebSocket();
|
||||
initChart();
|
||||
}).catch(err => {
|
||||
console.error('Initial load error:', err);
|
||||
addLogLine('Ошибка загрузки конфигурации: ' + err.message, 'warning');
|
||||
});
|
||||
});
|
||||
|
|
|
|||
5
src/yolo_webui/static/favicon.svg
Normal file
5
src/yolo_webui/static/favicon.svg
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="#f97316" stroke-width="2.2" stroke-linecap="round" stroke-linejoin="round">
|
||||
<path d="M21 16V8a2 2 0 0 0-1-1.73l-7-4a2 2 0 0 0-2 0l-7 4A2 2 0 0 0 3 8v8a2 2 0 0 0 1 1.73l7 4a2 2 0 0 0 2 0l7-4A2 2 0 0 0 21 16z"/>
|
||||
<polyline points="3.27 6.96 12 12.01 20.73 6.96"/>
|
||||
<line x1="12" y1="22.08" x2="12" y2="12"/>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 397 B |
|
|
@ -7,10 +7,8 @@
|
|||
<!-- Google Fonts -->
|
||||
<link rel="preconnect" href="https://fonts.googleapis.com">
|
||||
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
|
||||
<link href="https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap" rel="stylesheet">
|
||||
<!-- Chart.js -->
|
||||
<script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
|
||||
<link rel="stylesheet" href="/static/style.css">
|
||||
<link rel="icon" type="image/svg+xml" href="/static/favicon.svg">
|
||||
<link rel="stylesheet" href="/static/style.css?v=6">
|
||||
</head>
|
||||
<body>
|
||||
<header class="app-header">
|
||||
|
|
@ -30,7 +28,7 @@
|
|||
</div>
|
||||
</div>
|
||||
<div class="header-actions">
|
||||
<a href="http://localhost:5000" target="_blank" class="mlflow-link" id="mlflow-header-link">
|
||||
<a href="http://localhost:5000" target="_blank" rel="noopener noreferrer" class="mlflow-link" id="mlflow-header-link">
|
||||
<svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" class="link-icon">
|
||||
<path d="M18 13v6a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2V8a2 2 0 0 1 2-2h6"></path>
|
||||
<polyline points="15 3 21 3 21 9"></polyline>
|
||||
|
|
@ -70,7 +68,7 @@
|
|||
<div class="session-row">
|
||||
<div class="session-field">
|
||||
<label for="session-name">Сохранить текущие настройки как профиль</label>
|
||||
<input type="text" id="session-name" placeholder="Введите имя профиля...">
|
||||
<input type="text" id="session-name" pattern="[A-Za-z0-9_-]+" autocomplete="off" placeholder="Введите имя профиля...">
|
||||
</div>
|
||||
<button type="button" id="session-save-btn" class="btn-action btn-save" title="Сохранить настройки">
|
||||
<svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" class="btn-icon-small"><path d="M19 21H5a2 2 0 0 1-2-2V5a2 2 0 0 1 2-2h11l5 5v11a2 2 0 0 1-2 2z"></path><polyline points="17 21 17 13 7 13 7 21"></polyline></svg>
|
||||
|
|
@ -377,17 +375,6 @@
|
|||
</button>
|
||||
</div>
|
||||
|
||||
<!-- Live Chart -->
|
||||
<div class="chart-container-card">
|
||||
<div class="chart-header">
|
||||
<h3>График обучения (Live)</h3>
|
||||
<div class="chart-legend" id="chart-legend"></div>
|
||||
</div>
|
||||
<div class="chart-canvas-wrapper">
|
||||
<canvas id="metricsChart"></canvas>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Logs Box -->
|
||||
<div class="log-card">
|
||||
<div class="log-header">
|
||||
|
|
@ -433,11 +420,11 @@
|
|||
<option value="onnx">ONNX (.onnx)</option>
|
||||
<option value="engine">TensorRT (.engine)</option>
|
||||
<option value="openvino">OpenVINO (_openvino_model/)</option>
|
||||
<option value="triton">Triton Inference Server (triton/)</option>
|
||||
<option value="torchscript">TorchScript (.torchscript)</option>
|
||||
<option value="coreml">CoreML (.mlpackage)</option>
|
||||
<option value="tflite">TFLite (.tflite)</option>
|
||||
<option value="pb">TensorFlow SavedModel (_saved_model/)</option>
|
||||
<option value="saved_model">TensorFlow SavedModel (_saved_model/)</option>
|
||||
<option value="pb">TensorFlow GraphDef (.pb)</option>
|
||||
</select>
|
||||
</div>
|
||||
|
||||
|
|
@ -445,15 +432,15 @@
|
|||
<div class="row">
|
||||
<div class="field">
|
||||
<label for="export-imgsz">Размер изображения (imgsz)</label>
|
||||
<input type="number" id="export-imgsz" name="export-imgsz" value="640">
|
||||
<input type="number" id="export-imgsz" name="export-imgsz" value="640" min="32" max="8192" step="1" required>
|
||||
</div>
|
||||
<div class="field">
|
||||
<label for="export-batch" title="Максимальный размер батча. Важен при Dynamic=True для TensorRT и Triton.">Макс. батч (batch)</label>
|
||||
<input type="number" id="export-batch" name="export-batch" value="1" min="1">
|
||||
<label for="export-batch" title="Максимальный размер батча. Важен при Dynamic=True для TensorRT.">Макс. батч (batch)</label>
|
||||
<input type="number" id="export-batch" name="export-batch" value="1" min="1" max="1024" step="1" required>
|
||||
</div>
|
||||
<div class="field">
|
||||
<label for="export-workspace" title="Размер памяти для сборки TensorRT (ГБ).">Workspace (ГБ)</label>
|
||||
<input type="number" id="export-workspace" name="export-workspace" value="4" min="1" max="64">
|
||||
<input type="number" id="export-workspace" name="export-workspace" value="4" min="1" max="64" step="0.5" required>
|
||||
</div>
|
||||
</div>
|
||||
<div class="row">
|
||||
|
|
@ -549,6 +536,6 @@
|
|||
</section>
|
||||
</div> <!-- End export-view -->
|
||||
</main>
|
||||
<script src="/static/app.js"></script>
|
||||
<script src="/static/app.js?v=6"></script>
|
||||
</body>
|
||||
</html>
|
||||
|
|
|
|||
|
|
@ -44,10 +44,11 @@ body {
|
|||
background-color: var(--bg-primary);
|
||||
color: var(--text-main);
|
||||
font-family: var(--font-sans);
|
||||
min-height: 100vh;
|
||||
height: 100vh;
|
||||
max-height: 100vh;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
overflow-x: hidden;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
/* Header Styles */
|
||||
|
|
@ -136,14 +137,18 @@ body {
|
|||
width: 100%;
|
||||
margin: 0 auto;
|
||||
height: calc(100vh - 57px);
|
||||
min-height: 0;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.workspace-view {
|
||||
display: none;
|
||||
grid-template-columns: 46% 1fr;
|
||||
grid-template-columns: 420px 1fr;
|
||||
gap: 1.5rem;
|
||||
flex: 1;
|
||||
height: 100%;
|
||||
min-height: 0;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.workspace-view.active {
|
||||
|
|
@ -169,46 +174,58 @@ body {
|
|||
flex-direction: column;
|
||||
overflow: hidden;
|
||||
height: 100%;
|
||||
min-height: 0;
|
||||
}
|
||||
|
||||
#config-pane {
|
||||
#config-pane,
|
||||
#export-config-pane {
|
||||
padding: 1.25rem;
|
||||
max-height: 100%;
|
||||
}
|
||||
|
||||
#run-pane {
|
||||
#run-pane,
|
||||
#export-run-pane {
|
||||
padding: 1.25rem;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 1.25rem;
|
||||
height: 100%;
|
||||
max-height: 100%;
|
||||
overflow-y: auto;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
/* Custom Scrollbars */
|
||||
#run-pane::-webkit-scrollbar,
|
||||
#export-run-pane::-webkit-scrollbar,
|
||||
.form-container::-webkit-scrollbar,
|
||||
#log-container::-webkit-scrollbar {
|
||||
#log-container::-webkit-scrollbar,
|
||||
#export-log-container::-webkit-scrollbar {
|
||||
width: 6px;
|
||||
height: 6px;
|
||||
}
|
||||
|
||||
#run-pane::-webkit-scrollbar-track,
|
||||
#export-run-pane::-webkit-scrollbar-track,
|
||||
.form-container::-webkit-scrollbar-track,
|
||||
#log-container::-webkit-scrollbar-track {
|
||||
#log-container::-webkit-scrollbar-track,
|
||||
#export-log-container::-webkit-scrollbar-track {
|
||||
background: transparent;
|
||||
}
|
||||
|
||||
#run-pane::-webkit-scrollbar-thumb,
|
||||
#export-run-pane::-webkit-scrollbar-thumb,
|
||||
.form-container::-webkit-scrollbar-thumb,
|
||||
#log-container::-webkit-scrollbar-thumb {
|
||||
#log-container::-webkit-scrollbar-thumb,
|
||||
#export-log-container::-webkit-scrollbar-thumb {
|
||||
background: var(--border-color);
|
||||
border-radius: 3px;
|
||||
}
|
||||
|
||||
#run-pane::-webkit-scrollbar-thumb:hover,
|
||||
#export-run-pane::-webkit-scrollbar-thumb:hover,
|
||||
.form-container::-webkit-scrollbar-thumb:hover,
|
||||
#log-container::-webkit-scrollbar-thumb:hover {
|
||||
#log-container::-webkit-scrollbar-thumb:hover,
|
||||
#export-log-container::-webkit-scrollbar-thumb:hover {
|
||||
background: var(--border-hover);
|
||||
}
|
||||
|
||||
|
|
@ -333,20 +350,26 @@ body {
|
|||
border-radius: 8px;
|
||||
padding: 0.25rem;
|
||||
margin-bottom: 1.25rem;
|
||||
gap: 0.25rem;
|
||||
gap: 0.2rem;
|
||||
overflow-x: auto;
|
||||
}
|
||||
|
||||
.tab-btn {
|
||||
flex: 1;
|
||||
flex: 1 1 0;
|
||||
min-width: 0;
|
||||
background: none;
|
||||
border: none;
|
||||
border-radius: 6px;
|
||||
color: var(--text-muted);
|
||||
font-family: var(--font-sans);
|
||||
font-size: 0.9rem;
|
||||
font-size: 0.82rem;
|
||||
font-weight: 600;
|
||||
padding: 0.6rem;
|
||||
padding: 0.55rem 0.25rem;
|
||||
cursor: pointer;
|
||||
text-align: center;
|
||||
white-space: nowrap;
|
||||
overflow: hidden;
|
||||
text-overflow: ellipsis;
|
||||
transition: var(--transition-fast);
|
||||
}
|
||||
|
||||
|
|
@ -543,7 +566,8 @@ body {
|
|||
}
|
||||
|
||||
/* Status Cards & Themes */
|
||||
#status-card {
|
||||
#status-card,
|
||||
#export-status-card {
|
||||
border-radius: 8px;
|
||||
padding: 1.25rem;
|
||||
border-left: 5px solid var(--text-dim);
|
||||
|
|
@ -602,7 +626,8 @@ body {
|
|||
100% { transform: scale(0.9); opacity: 0.6; }
|
||||
}
|
||||
|
||||
#status-title {
|
||||
#status-title,
|
||||
#export-status-title {
|
||||
font-size: 1rem;
|
||||
font-weight: 700;
|
||||
letter-spacing: 0.05em;
|
||||
|
|
@ -624,12 +649,33 @@ body {
|
|||
color: var(--text-muted);
|
||||
}
|
||||
|
||||
#status-text {
|
||||
#status-text,
|
||||
#export-status-text {
|
||||
font-size: 0.9rem;
|
||||
color: var(--text-main);
|
||||
margin-bottom: 1rem;
|
||||
}
|
||||
|
||||
#export-status-card.status-preparing,
|
||||
#export-status-card.status-stopping { border-left-color: var(--warning); }
|
||||
#export-status-card.status-exporting,
|
||||
#export-status-card.status-succeeded { border-left-color: var(--success); }
|
||||
#export-status-card.status-cancelled { border-left-color: var(--warning); }
|
||||
#export-status-card.status-failed { border-left-color: var(--error); }
|
||||
#export-status-card.status-preparing .status-dot,
|
||||
#export-status-card.status-stopping .status-dot { background-color: var(--warning); box-shadow: 0 0 8px var(--warning); }
|
||||
#export-status-card.status-exporting .status-dot,
|
||||
#export-status-card.status-succeeded .status-dot { background-color: var(--success); box-shadow: 0 0 8px var(--success); }
|
||||
#export-status-card.status-cancelled .status-dot { background-color: var(--warning); box-shadow: 0 0 8px var(--warning); }
|
||||
#export-status-card.status-failed .status-dot { background-color: var(--error); box-shadow: 0 0 8px var(--error); }
|
||||
#export-status-card.status-preparing #export-status-title,
|
||||
#export-status-card.status-stopping #export-status-title { color: var(--warning); }
|
||||
#export-status-card.status-exporting #export-status-title,
|
||||
#export-status-card.status-succeeded #export-status-title { color: var(--success); }
|
||||
#export-status-card.status-cancelled #export-status-title { color: var(--warning); }
|
||||
#export-status-card.status-failed #export-status-title { color: var(--error); }
|
||||
|
||||
|
||||
/* Progress bar inside status card */
|
||||
.progress-container {
|
||||
display: flex;
|
||||
|
|
@ -721,40 +767,6 @@ body {
|
|||
box-shadow: none !important;
|
||||
}
|
||||
|
||||
/* Charts Card */
|
||||
.chart-container-card {
|
||||
background-color: var(--bg-secondary);
|
||||
border: 1px solid var(--border-color);
|
||||
border-radius: 8px;
|
||||
padding: 1rem;
|
||||
height: 320px;
|
||||
min-height: 320px;
|
||||
flex: none;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
}
|
||||
|
||||
.chart-header {
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
align-items: center;
|
||||
margin-bottom: 0.75rem;
|
||||
}
|
||||
|
||||
.chart-header h3 {
|
||||
font-size: 0.9rem;
|
||||
font-weight: 700;
|
||||
text-transform: uppercase;
|
||||
letter-spacing: 0.05em;
|
||||
color: var(--text-muted);
|
||||
}
|
||||
|
||||
.chart-canvas-wrapper {
|
||||
position: relative;
|
||||
width: 100%;
|
||||
height: 240px;
|
||||
}
|
||||
|
||||
/* Console Logs Box */
|
||||
.log-card {
|
||||
background-color: #0d0d0f;
|
||||
|
|
@ -762,9 +774,8 @@ body {
|
|||
border-radius: 8px;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
height: 280px;
|
||||
min-height: 280px;
|
||||
flex: none;
|
||||
flex: 1;
|
||||
min-height: 0;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
|
|
@ -811,21 +822,24 @@ body {
|
|||
}
|
||||
|
||||
.log-body {
|
||||
flex: 1;
|
||||
flex: 1 1 0;
|
||||
min-height: 0;
|
||||
overflow-x: auto;
|
||||
overflow-y: auto;
|
||||
padding: 0.75rem 1rem;
|
||||
padding: 0.6rem 0.85rem;
|
||||
font-family: var(--font-mono);
|
||||
font-size: 0.85rem;
|
||||
line-height: 1.5;
|
||||
font-size: 0.75rem;
|
||||
line-height: 1.35;
|
||||
letter-spacing: -0.015em;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 0.25rem;
|
||||
gap: 0.15rem;
|
||||
color: #e4e4e7;
|
||||
}
|
||||
|
||||
.log-line {
|
||||
white-space: pre-wrap;
|
||||
word-break: break-all;
|
||||
white-space: pre;
|
||||
word-break: normal;
|
||||
}
|
||||
|
||||
.log-level-info { color: var(--text-muted); }
|
||||
|
|
|
|||
|
|
@ -189,26 +189,43 @@ class TrainingRunner:
|
|||
|
||||
on_event(TrainingEvent("info", "Загрузка Ultralytics и подготовка модели…"))
|
||||
from ultralytics import YOLO, settings
|
||||
from .ultralytics_trainers import trainer_for_task
|
||||
|
||||
settings.update({"mlflow": config.mlflow.enabled})
|
||||
workspace_root = Path.cwd().resolve()
|
||||
try:
|
||||
settings.update({
|
||||
"runs_dir": str((workspace_root / "runs").resolve()),
|
||||
"datasets_dir": str((workspace_root / "datasets").resolve()),
|
||||
"weights_dir": str((workspace_root / "models").resolve()),
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
with mlflow_environment(config.mlflow):
|
||||
model = YOLO(config.resolved_model, task=config.task)
|
||||
with self._state_lock:
|
||||
self._model = model
|
||||
|
||||
model.add_callback("on_train_start", self._on_train_start(on_event))
|
||||
model.add_callback("on_train_epoch_end", self._on_epoch_end(on_event))
|
||||
model.add_callback("on_train_end", self._on_train_end(on_event))
|
||||
|
||||
try:
|
||||
model.train(**train_args)
|
||||
trainer = getattr(model, "trainer", None)
|
||||
save_dir = getattr(trainer, "save_dir", None)
|
||||
return Path(save_dir) if save_dir else None
|
||||
finally:
|
||||
previous_mlflow_setting = settings["mlflow"]
|
||||
setting_changed = previous_mlflow_setting is not config.mlflow.enabled
|
||||
if setting_changed:
|
||||
settings.update({"mlflow": config.mlflow.enabled})
|
||||
try:
|
||||
with mlflow_environment(config.mlflow):
|
||||
model = YOLO(config.resolved_model, task=config.task)
|
||||
with self._state_lock:
|
||||
self._model = None
|
||||
self._model = model
|
||||
|
||||
model.add_callback("on_train_start", self._on_train_start(on_event))
|
||||
model.add_callback("on_fit_epoch_end", self._on_epoch_end(on_event))
|
||||
model.add_callback("on_train_end", self._on_train_end(on_event))
|
||||
|
||||
try:
|
||||
model.train(trainer=trainer_for_task(config.task), **train_args)
|
||||
trainer = getattr(model, "trainer", None)
|
||||
save_dir = getattr(trainer, "save_dir", None)
|
||||
return Path(save_dir) if save_dir else None
|
||||
finally:
|
||||
with self._state_lock:
|
||||
self._model = None
|
||||
finally:
|
||||
if setting_changed:
|
||||
settings.update({"mlflow": previous_mlflow_setting})
|
||||
|
||||
def _on_train_start(self, on_event: EventHandler) -> Callable[[Any], None]:
|
||||
def callback(trainer: Any) -> None:
|
||||
|
|
@ -221,6 +238,9 @@ class TrainingRunner:
|
|||
|
||||
def _on_epoch_end(self, on_event: EventHandler) -> Callable[[Any], None]:
|
||||
def callback(trainer: Any) -> None:
|
||||
validator = getattr(trainer, "validator", None)
|
||||
if validator is not None and getattr(validator, "training", True) is False:
|
||||
return # final best-checkpoint validation is not an additional train epoch
|
||||
epoch = int(getattr(trainer, "epoch", 0)) + 1
|
||||
total = int(getattr(getattr(trainer, "args", None), "epochs", 0))
|
||||
metrics = getattr(trainer, "metrics", {}) or {}
|
||||
|
|
|
|||
72
src/yolo_webui/ultralytics_trainers.py
Normal file
72
src/yolo_webui/ultralytics_trainers.py
Normal file
|
|
@ -0,0 +1,72 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import Any, ClassVar
|
||||
|
||||
from ultralytics.models.yolo.classify.train import ClassificationTrainer
|
||||
from ultralytics.models.yolo.detect.train import DetectionTrainer
|
||||
from ultralytics.models.yolo.obb.train import OBBTrainer
|
||||
from ultralytics.models.yolo.pose.train import PoseTrainer
|
||||
from ultralytics.models.yolo.segment.train import SegmentationTrainer
|
||||
from ultralytics.utils import SETTINGS
|
||||
|
||||
from .config import YoloTask
|
||||
from .mlflow_metrics import TaskMetricsMonitor
|
||||
|
||||
|
||||
class _TaskMetricsTrainerMixin:
|
||||
"""Install monitoring inside the trainer so Ultralytics DDP keeps it."""
|
||||
|
||||
monitoring_task: ClassVar[YoloTask]
|
||||
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self._install_task_metrics_monitor()
|
||||
|
||||
def _install_task_metrics_monitor(self) -> None:
|
||||
monitor = TaskMetricsMonitor(
|
||||
self.monitoring_task,
|
||||
mlflow_enabled=SETTINGS["mlflow"] is True,
|
||||
)
|
||||
self._task_metrics_monitor = monitor
|
||||
callbacks = (
|
||||
("on_train_start", monitor.on_train_start),
|
||||
("on_fit_epoch_end", monitor.on_fit_epoch_end),
|
||||
("on_train_end", monitor.on_train_end),
|
||||
)
|
||||
for event, callback in callbacks:
|
||||
# BaseTrainer adds integrations in __init__. Prepending ensures our
|
||||
# metrics and CSV exist before Ultralytics logs them to MLflow.
|
||||
self.callbacks.setdefault(event, []).insert(0, callback)
|
||||
|
||||
|
||||
class MonitoredDetectionTrainer(_TaskMetricsTrainerMixin, DetectionTrainer):
|
||||
monitoring_task = "detect"
|
||||
|
||||
|
||||
class MonitoredSegmentationTrainer(_TaskMetricsTrainerMixin, SegmentationTrainer):
|
||||
monitoring_task = "segment"
|
||||
|
||||
|
||||
class MonitoredClassificationTrainer(_TaskMetricsTrainerMixin, ClassificationTrainer):
|
||||
monitoring_task = "classify"
|
||||
|
||||
|
||||
class MonitoredPoseTrainer(_TaskMetricsTrainerMixin, PoseTrainer):
|
||||
monitoring_task = "pose"
|
||||
|
||||
|
||||
class MonitoredOBBTrainer(_TaskMetricsTrainerMixin, OBBTrainer):
|
||||
monitoring_task = "obb"
|
||||
|
||||
|
||||
MONITORED_TRAINERS = {
|
||||
"detect": MonitoredDetectionTrainer,
|
||||
"segment": MonitoredSegmentationTrainer,
|
||||
"classify": MonitoredClassificationTrainer,
|
||||
"pose": MonitoredPoseTrainer,
|
||||
"obb": MonitoredOBBTrainer,
|
||||
}
|
||||
|
||||
|
||||
def trainer_for_task(task: YoloTask) -> type:
|
||||
return MONITORED_TRAINERS[task]
|
||||
|
|
@ -76,21 +76,27 @@ global.localStorage = {
|
|||
removeItem(key) { storage.delete(key); }
|
||||
};
|
||||
global.confirm = () => true;
|
||||
global.window = {location: {protocol: 'http:', host: '127.0.0.1:8000'}};
|
||||
|
||||
class FakeChart {
|
||||
static instances = [];
|
||||
|
||||
constructor(_context, config) {
|
||||
this.data = config.data;
|
||||
this.options = config.options;
|
||||
FakeChart.instances.push(this);
|
||||
global.window = {
|
||||
location: {
|
||||
protocol: 'http:',
|
||||
host: '127.0.0.1:8000',
|
||||
hostname: '127.0.0.1'
|
||||
}
|
||||
};
|
||||
|
||||
let nextTimerId = 1;
|
||||
const scheduledTimers = [];
|
||||
global.setTimeout = (handler, delay) => {
|
||||
const timer = {id: nextTimerId++, handler, delay, cancelled: false, fired: false};
|
||||
scheduledTimers.push(timer);
|
||||
return timer.id;
|
||||
};
|
||||
global.clearTimeout = id => {
|
||||
const timer = scheduledTimers.find(item => item.id === id);
|
||||
if (timer) timer.cancelled = true;
|
||||
};
|
||||
|
||||
|
||||
destroy() {}
|
||||
update() {}
|
||||
}
|
||||
global.Chart = FakeChart;
|
||||
|
||||
class FakeWebSocket {
|
||||
static instances = [];
|
||||
|
|
@ -106,8 +112,25 @@ const response = (ok, data) => ({
|
|||
ok,
|
||||
async json() { return data; }
|
||||
});
|
||||
global.fetch = async url => {
|
||||
const fetchCalls = [];
|
||||
let holdTrainingStart = false;
|
||||
let holdExportStart = false;
|
||||
global.fetch = async (url, options = {}) => {
|
||||
fetchCalls.push({url, options});
|
||||
if (url === '/api/train/start' && holdTrainingStart) {
|
||||
return new Promise(() => {});
|
||||
}
|
||||
if (url === '/api/export/start' && holdExportStart) {
|
||||
return new Promise(() => {});
|
||||
}
|
||||
if (url === '/api/datasets' || url === '/api/models' || url === '/api/sessions') {
|
||||
if (url === '/api/models') {
|
||||
return response(true, [{
|
||||
name: 'detect/experiment/best.pt',
|
||||
path: '/workspace/runs/detect/experiment/weights/best.pt',
|
||||
source: 'runs'
|
||||
}]);
|
||||
}
|
||||
return response(true, []);
|
||||
}
|
||||
if (url === '/api/sessions/last_run') return response(false, {});
|
||||
|
|
@ -119,7 +142,7 @@ global.fetch = async url => {
|
|||
workers: 8,
|
||||
patience: 100,
|
||||
augmentation: {enabled: true, close_mosaic: 10},
|
||||
mlflow: {enabled: false},
|
||||
mlflow: {enabled: false, tracking_uri: 'http://mlflow:5000'},
|
||||
split: {enabled: false}
|
||||
});
|
||||
}
|
||||
|
|
@ -133,11 +156,30 @@ async function flushPromises() {
|
|||
await new Promise(resolve => setImmediate(resolve));
|
||||
}
|
||||
|
||||
function descendants(node) {
|
||||
return node.children.flatMap(child => [child, ...descendants(child)]);
|
||||
}
|
||||
|
||||
function runTimer(timer) {
|
||||
timer.fired = true;
|
||||
timer.handler();
|
||||
}
|
||||
|
||||
(async () => {
|
||||
assert.equal(typeof domReady, 'function');
|
||||
domReady();
|
||||
await flushPromises();
|
||||
|
||||
assert.equal(element('mlflow-header-link').href, 'http://127.0.0.1:5000/');
|
||||
|
||||
const localModelPath = '/workspace/runs/detect/experiment/weights/best.pt';
|
||||
const localModelOption = descendants(element('model-select'))
|
||||
.find(item => item.textContent === 'detect/experiment/best.pt');
|
||||
assert.ok(localModelOption);
|
||||
assert.equal(localModelOption.value, localModelPath);
|
||||
|
||||
element('model-select').value = localModelOption.value;
|
||||
element('model-select').listeners.change();
|
||||
element('workers').value = '0';
|
||||
element('patience').value = '0';
|
||||
element('close-mosaic').value = '0';
|
||||
|
|
@ -148,9 +190,28 @@ async function flushPromises() {
|
|||
assert.equal(savedConfig.patience, 0);
|
||||
assert.equal(savedConfig.augmentation.close_mosaic, 0);
|
||||
assert.equal(savedConfig.augmentation.auto_augment, 'none');
|
||||
assert.equal(savedConfig.model, localModelPath);
|
||||
|
||||
element('model-select').value = '__custom__';
|
||||
element('model-select').listeners.change();
|
||||
element('model').value = '/workspace/models/custom-segment.pt';
|
||||
element('task').value = 'segment';
|
||||
element('task').listeners.change();
|
||||
assert.equal(element('model-select').value, '__custom__');
|
||||
assert.equal(element('model').value, '/workspace/models/custom-segment.pt');
|
||||
|
||||
assert.equal(FakeWebSocket.instances.length, 2);
|
||||
const socket = FakeWebSocket.instances[0];
|
||||
const exportSocket = FakeWebSocket.instances[1];
|
||||
const originalConsoleError = console.error;
|
||||
console.error = () => {};
|
||||
try {
|
||||
assert.doesNotThrow(() => socket.onmessage({data: '{broken'}));
|
||||
assert.doesNotThrow(() => exportSocket.onmessage({data: 'null'}));
|
||||
} finally {
|
||||
console.error = originalConsoleError;
|
||||
}
|
||||
|
||||
socket.onmessage({
|
||||
data: JSON.stringify({
|
||||
type: 'init',
|
||||
|
|
@ -163,11 +224,76 @@ async function flushPromises() {
|
|||
})
|
||||
});
|
||||
|
||||
const chart = FakeChart.instances.at(-1);
|
||||
assert.deepEqual(chart.data.labels, [1, 2]);
|
||||
assert.deepEqual(chart.data.datasets.map(item => item.label), ['mAP50', 'loss']);
|
||||
assert.deepEqual(chart.data.datasets[0].data, [0.5, null]);
|
||||
assert.deepEqual(chart.data.datasets[1].data, [null, 0.2]);
|
||||
socket.onclose();
|
||||
socket.onclose();
|
||||
let activeReconnectTimers = scheduledTimers.filter(
|
||||
timer => timer.delay === 5000 && !timer.cancelled && !timer.fired
|
||||
);
|
||||
assert.equal(activeReconnectTimers.length, 1);
|
||||
runTimer(activeReconnectTimers[0]);
|
||||
assert.equal(FakeWebSocket.instances.length, 3);
|
||||
assert.equal(FakeWebSocket.instances[2].url, 'ws://127.0.0.1:8000/api/ws');
|
||||
|
||||
exportSocket.onclose();
|
||||
exportSocket.onclose();
|
||||
activeReconnectTimers = scheduledTimers.filter(
|
||||
timer => timer.delay === 5000 && !timer.cancelled && !timer.fired
|
||||
);
|
||||
assert.equal(activeReconnectTimers.length, 1);
|
||||
runTimer(activeReconnectTimers[0]);
|
||||
assert.equal(FakeWebSocket.instances.length, 4);
|
||||
assert.equal(FakeWebSocket.instances[3].url, 'ws://127.0.0.1:8000/api/export/ws');
|
||||
|
||||
element('session-name').value = 'profile with spaces';
|
||||
await element('session-save-btn').listeners.click();
|
||||
assert.equal(
|
||||
fetchCalls.filter(call => call.options.method === 'POST' && call.url.startsWith('/api/sessions/')).length,
|
||||
0
|
||||
);
|
||||
|
||||
element('export-model').value = localModelPath;
|
||||
element('export-format').value = 'pb';
|
||||
element('export-imgsz').value = '31';
|
||||
element('export-batch').value = '1';
|
||||
element('export-workspace').value = '4';
|
||||
await element('export-start-btn').listeners.click();
|
||||
assert.equal(fetchCalls.filter(call => call.url === '/api/export/start').length, 0);
|
||||
|
||||
element('export-imgsz').value = '640';
|
||||
element('export-workspace').value = '1.5';
|
||||
element('export-half').checked = true;
|
||||
element('export-int8').checked = true;
|
||||
await element('export-start-btn').listeners.click();
|
||||
assert.equal(fetchCalls.filter(call => call.url === '/api/export/start').length, 0);
|
||||
|
||||
element('export-half').checked = false;
|
||||
element('export-int8').checked = false;
|
||||
holdExportStart = true;
|
||||
element('export-start-btn').listeners.click();
|
||||
element('export-start-btn').listeners.click();
|
||||
const exportStartCalls = fetchCalls.filter(call => call.url === '/api/export/start');
|
||||
assert.equal(exportStartCalls.length, 1);
|
||||
assert.equal(element('export-start-btn').disabled, true);
|
||||
assert.deepEqual(
|
||||
JSON.parse(exportStartCalls[0].options.body),
|
||||
{
|
||||
model: localModelPath,
|
||||
format: 'pb',
|
||||
imgsz: 640,
|
||||
half: false,
|
||||
int8: false,
|
||||
dynamic: false,
|
||||
simplify: false,
|
||||
batch: 1,
|
||||
workspace: 1.5
|
||||
}
|
||||
);
|
||||
|
||||
holdTrainingStart = true;
|
||||
element('start-btn').listeners.click();
|
||||
element('start-btn').listeners.click();
|
||||
assert.equal(fetchCalls.filter(call => call.url === '/api/train/start').length, 1);
|
||||
assert.equal(element('start-btn').disabled, true);
|
||||
})().catch(error => {
|
||||
console.error(error);
|
||||
process.exitCode = 1;
|
||||
|
|
|
|||
|
|
@ -5,8 +5,16 @@ import json
|
|||
import threading
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
import pytest
|
||||
|
||||
from yolo_webui.app import TrainingManager, app
|
||||
import yolo_webui.app as app_module
|
||||
from yolo_webui.app import (
|
||||
ExportManager,
|
||||
LiveState,
|
||||
TrainingManager,
|
||||
app,
|
||||
write_json_atomic,
|
||||
)
|
||||
|
||||
|
||||
def test_get_config_defaults() -> None:
|
||||
|
|
@ -94,6 +102,28 @@ def test_sessions_flow(monkeypatch, tmp_path) -> None:
|
|||
assert response.status_code == 404
|
||||
|
||||
|
||||
def test_invalid_session_config_is_not_saved(monkeypatch, tmp_path) -> None:
|
||||
client = TestClient(app)
|
||||
monkeypatch.setattr("yolo_webui.app.get_sessions_dir", lambda: tmp_path)
|
||||
|
||||
response = client.post("/api/sessions/broken", json={"unexpected": True})
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "датасета" in response.json()["detail"].lower()
|
||||
assert not (tmp_path / "broken.json").exists()
|
||||
|
||||
|
||||
def test_atomic_json_write_cleans_temp_file_after_serialization_error(tmp_path) -> None:
|
||||
destination = tmp_path / "session.json"
|
||||
destination.write_text('{"old": true}', encoding="utf-8")
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
write_json_atomic(destination, {"invalid": object()})
|
||||
|
||||
assert destination.read_text(encoding="utf-8") == '{"old": true}'
|
||||
assert list(tmp_path.iterdir()) == [destination]
|
||||
|
||||
|
||||
def test_reserved_session_name_cannot_be_overwritten(monkeypatch, tmp_path) -> None:
|
||||
client = TestClient(app)
|
||||
monkeypatch.setattr("yolo_webui.app.get_sessions_dir", lambda: tmp_path)
|
||||
|
|
@ -104,6 +134,107 @@ def test_reserved_session_name_cannot_be_overwritten(monkeypatch, tmp_path) -> N
|
|||
assert "зарезервировано" in response.json()["detail"]
|
||||
|
||||
|
||||
def test_rejected_busy_start_does_not_overwrite_last_run(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
) -> None:
|
||||
client = TestClient(app)
|
||||
previous = {"dataset": "previous.yaml", "model": "previous.pt"}
|
||||
last_run_path = tmp_path / "last_run.json"
|
||||
last_run_path.write_text(json.dumps(previous), encoding="utf-8")
|
||||
monkeypatch.setattr("yolo_webui.app.get_sessions_dir", lambda: tmp_path)
|
||||
|
||||
def reject_start(_config) -> None:
|
||||
raise ValueError("Обучение уже выполняется.")
|
||||
|
||||
monkeypatch.setattr(app_module.manager, "start_training", reject_start)
|
||||
|
||||
response = client.post(
|
||||
"/api/train/start",
|
||||
json={"dataset": "coco8.yaml", "model": "yolo11n.pt"},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert json.loads(last_run_path.read_text(encoding="utf-8")) == previous
|
||||
|
||||
|
||||
def test_start_api_rejects_non_finite_numeric_value() -> None:
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post(
|
||||
"/api/train/start",
|
||||
content='{"dataset":"coco8.yaml","model":"yolo11n.pt","epochs":NaN}',
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "целым числом" in response.json()["detail"]
|
||||
|
||||
|
||||
def test_list_models_uses_path_relative_to_runs(monkeypatch, tmp_path) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
direct_model = tmp_path / "runs" / "direct.pt"
|
||||
nested_model = tmp_path / "runs" / "detect" / "train" / "weights" / "best.pt"
|
||||
direct_model.parent.mkdir()
|
||||
nested_model.parent.mkdir(parents=True)
|
||||
direct_model.write_bytes(b"")
|
||||
nested_model.write_bytes(b"")
|
||||
|
||||
response = TestClient(app).get("/api/models")
|
||||
|
||||
assert response.status_code == 200
|
||||
run_names = {
|
||||
item["name"] for item in response.json() if item["source"] == "runs"
|
||||
}
|
||||
assert run_names == {"direct.pt", "detect/train/weights/best.pt"}
|
||||
|
||||
|
||||
def test_export_start_rejects_invalid_config_synchronously() -> None:
|
||||
response = TestClient(app).post(
|
||||
"/api/export/start",
|
||||
json={"model": "missing.pt", "format": "not-a-format"},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "Неподдерживаемый формат" in response.json()["detail"]
|
||||
|
||||
|
||||
def test_export_start_passes_canonical_validated_config(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
model_path = tmp_path / "models" / "model.pt"
|
||||
model_path.parent.mkdir()
|
||||
model_path.write_bytes(b"checkpoint")
|
||||
received: list[dict[str, object]] = []
|
||||
monkeypatch.setattr(app_module.export_manager, "start_export", received.append)
|
||||
|
||||
response = TestClient(app).post(
|
||||
"/api/export/start",
|
||||
json={
|
||||
"model": "model.pt",
|
||||
"format": "ONNX",
|
||||
"workspace": 2.5,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert received == [
|
||||
{
|
||||
"model": str(model_path),
|
||||
"format": "onnx",
|
||||
"imgsz": 640,
|
||||
"half": False,
|
||||
"int8": False,
|
||||
"dynamic": False,
|
||||
"simplify": False,
|
||||
"batch": 1,
|
||||
"workspace": 2.5,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_started_event_does_not_deadlock() -> None:
|
||||
training_manager = TrainingManager()
|
||||
training_manager.state.status = "preparing"
|
||||
|
|
@ -173,3 +304,55 @@ def test_process_result_distinguishes_success_cancellation_and_failure() -> None
|
|||
failed_after_stop.state.stop_requested = True
|
||||
failed_after_stop._finalize_process_result(1)
|
||||
assert failed_after_stop.state.status == "failed"
|
||||
|
||||
|
||||
def test_live_state_snapshot_is_detached() -> None:
|
||||
state = LiveState(
|
||||
status="training",
|
||||
logs=["first"],
|
||||
metrics=[{"epoch": 1, "loss": 0.5}],
|
||||
)
|
||||
|
||||
snapshot = state.snapshot()
|
||||
state.logs.append("second")
|
||||
state.metrics[0]["loss"] = 0.25
|
||||
|
||||
assert snapshot["logs"] == ["first"]
|
||||
assert snapshot["metrics"] == [{"epoch": 1, "loss": 0.5}]
|
||||
|
||||
|
||||
def test_export_stop_does_not_relabel_completed_failure_as_cancelled() -> None:
|
||||
class CompletedProcess:
|
||||
returncode = 1
|
||||
|
||||
def poll(self) -> int:
|
||||
return self.returncode
|
||||
|
||||
def terminate(self) -> None:
|
||||
raise AssertionError("completed process must not be terminated")
|
||||
|
||||
export = ExportManager()
|
||||
export.state.status = "exporting"
|
||||
export._process = CompletedProcess() # type: ignore[assignment]
|
||||
|
||||
export.stop_export()
|
||||
export._finalize_process_result(1)
|
||||
|
||||
assert export.state.stop_requested is False
|
||||
assert export.state.status == "failed"
|
||||
|
||||
|
||||
def test_export_stopped_while_preparing_never_starts_process(monkeypatch) -> None:
|
||||
export = ExportManager()
|
||||
export.state.status = "preparing"
|
||||
export.stop_export()
|
||||
|
||||
def unexpected_popen(*_args, **_kwargs):
|
||||
raise AssertionError("subprocess must not start after cancellation")
|
||||
|
||||
monkeypatch.setattr(app_module.subprocess, "Popen", unexpected_popen)
|
||||
|
||||
export._run_subprocess({"model": "unused.pt", "format": "onnx"})
|
||||
|
||||
assert export.state.status == "cancelled"
|
||||
assert export._process is None
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ def test_train_kwargs_omit_optional_empty_values() -> None:
|
|||
augmentation=AugmentationConfig(enabled=False),
|
||||
)
|
||||
|
||||
from pathlib import Path
|
||||
assert config.train_kwargs() == {
|
||||
"data": "dataset.yaml",
|
||||
"epochs": 100,
|
||||
|
|
@ -22,7 +23,7 @@ def test_train_kwargs_omit_optional_empty_values() -> None:
|
|||
"batch": 16,
|
||||
"workers": 8,
|
||||
"patience": 100,
|
||||
"project": "runs/train",
|
||||
"project": str(Path("runs/train").resolve()),
|
||||
"verbose": True,
|
||||
}
|
||||
|
||||
|
|
@ -72,14 +73,102 @@ def test_augmentation_probabilities_are_validated(field: str) -> None:
|
|||
augmentation.validate()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch", [0, -2])
|
||||
def test_invalid_batch_is_rejected(batch: int) -> None:
|
||||
@pytest.mark.parametrize("batch", [0, -2, -0.5, 1.5])
|
||||
def test_invalid_batch_is_rejected(batch: int | float) -> None:
|
||||
config = TrainingConfig(dataset="dataset.yaml", model="model.pt", batch_size=batch)
|
||||
|
||||
with pytest.raises(ValueError, match="Batch"):
|
||||
config.validate()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch", [-1, 1, 16, 0.25])
|
||||
def test_supported_batch_modes_are_accepted(batch: int | float) -> None:
|
||||
TrainingConfig(
|
||||
dataset="dataset.yaml",
|
||||
model="model.pt",
|
||||
batch_size=batch,
|
||||
).validate()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value", "message"),
|
||||
[
|
||||
("epochs", float("nan"), "целым числом"),
|
||||
("epochs", True, "целым числом"),
|
||||
("workers", 1.5, "целым числом"),
|
||||
("patience", float("inf"), "целым числом"),
|
||||
],
|
||||
)
|
||||
def test_training_numeric_fields_reject_wrong_types_and_non_finite_values(
|
||||
field: str,
|
||||
value: object,
|
||||
message: str,
|
||||
) -> None:
|
||||
values = {"dataset": "dataset.yaml", "model": "model.pt", field: value}
|
||||
|
||||
with pytest.raises(ValueError, match=message):
|
||||
TrainingConfig(**values).validate()
|
||||
|
||||
|
||||
def test_augmentation_rejects_non_finite_unbounded_values() -> None:
|
||||
config = TrainingConfig(
|
||||
dataset="dataset.yaml",
|
||||
model="model.pt",
|
||||
augmentation=AugmentationConfig(degrees=float("inf")),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="конечным числом"):
|
||||
config.validate()
|
||||
|
||||
|
||||
def test_augmentation_rejects_integer_too_large_for_float() -> None:
|
||||
config = TrainingConfig(
|
||||
dataset="dataset.yaml",
|
||||
model="model.pt",
|
||||
augmentation=AugmentationConfig(degrees=10**10_000),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="конечным числом"):
|
||||
config.validate()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"run_name",
|
||||
["../escape", "/tmp/escape", r"..\\escape", "C:escape"],
|
||||
)
|
||||
def test_run_name_cannot_escape_project(run_name: str) -> None:
|
||||
config = TrainingConfig(
|
||||
dataset="dataset.yaml",
|
||||
model="model.pt",
|
||||
run_name=run_name,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Имя запуска"):
|
||||
config.validate()
|
||||
|
||||
|
||||
def test_nested_config_sections_must_be_objects() -> None:
|
||||
with pytest.raises(ValueError, match="разделения датасета.*JSON-объектом"):
|
||||
TrainingConfig.from_dict(
|
||||
{
|
||||
"dataset": "dataset.yaml",
|
||||
"model": "model.pt",
|
||||
"split": None,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_classes_path_type_is_validated_before_path_operations() -> None:
|
||||
config = TrainingConfig(
|
||||
dataset="dataset.yaml",
|
||||
model="model.pt",
|
||||
split=DatasetSplitConfig(classes_path=123), # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Путь к файлу классов должен быть строкой"):
|
||||
config.validate()
|
||||
|
||||
|
||||
def test_mlflow_environment_is_restored(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("MLFLOW_TRACKING_URI", "previous")
|
||||
config = MlflowConfig(
|
||||
|
|
@ -95,6 +184,14 @@ def test_mlflow_environment_is_restored(monkeypatch: pytest.MonkeyPatch) -> None
|
|||
assert os.environ["MLFLOW_RUN"] == config.run_name
|
||||
|
||||
assert os.environ["MLFLOW_TRACKING_URI"] == "previous"
|
||||
|
||||
|
||||
def test_mlflow_default_tracking_uri_can_be_overridden_by_environment(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("YOLO_WEBUI_MLFLOW_TRACKING_URI", "http://mlflow:5000")
|
||||
|
||||
assert MlflowConfig().tracking_uri == "http://mlflow:5000"
|
||||
assert "MLFLOW_EXPERIMENT_NAME" not in os.environ
|
||||
assert "MLFLOW_RUN" not in os.environ
|
||||
|
||||
|
|
@ -113,7 +210,7 @@ def test_metrics_summary_skips_non_numeric_values() -> None:
|
|||
assert summary == "mAP50=0.8123 · loss=0.1235 · precision=0.9568"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ratio", [0.05, 0.98])
|
||||
@pytest.mark.parametrize("ratio", [0.05, 0.98, float("nan")])
|
||||
def test_dataset_split_ratio_is_validated(ratio: float) -> None:
|
||||
config = DatasetSplitConfig(enabled=True, train_ratio=ratio)
|
||||
with pytest.raises(ValueError, match="Доля обучающей выборки"):
|
||||
|
|
|
|||
41
tests/test_deployment.py
Normal file
41
tests/test_deployment.py
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def test_compose_is_local_and_does_not_require_nvidia() -> None:
|
||||
compose = yaml.safe_load(
|
||||
(PROJECT_ROOT / "docker-compose.yml").read_text(encoding="utf-8")
|
||||
)
|
||||
assert set(compose["services"].keys()) == {"webui"}
|
||||
webui = compose["services"]["webui"]
|
||||
|
||||
assert webui["ports"] == ["8000:8000"]
|
||||
assert "deploy" not in webui
|
||||
assert "YOLO_WEBUI_MLFLOW_TRACKING_URI" in webui["environment"]
|
||||
|
||||
|
||||
def test_docker_dependencies_are_exported_from_lock_file() -> None:
|
||||
dockerfile = (PROJECT_ROOT / "Dockerfile").read_text(encoding="utf-8")
|
||||
|
||||
assert "COPY pyproject.toml uv.lock README.md ./" in dockerfile
|
||||
assert "uv export --locked --no-dev" in dockerfile
|
||||
assert "--group build" in dockerfile
|
||||
assert "--prune ultralytics" in dockerfile
|
||||
assert "uv pip install --system --no-deps --no-build-isolation ." in dockerfile
|
||||
assert "ULTRALYTICS_SAFE_LOAD=1" in dockerfile
|
||||
|
||||
|
||||
def test_docker_context_excludes_local_training_artifacts() -> None:
|
||||
ignored = {
|
||||
line.strip()
|
||||
for line in (PROJECT_ROOT / ".dockerignore").read_text(encoding="utf-8").splitlines()
|
||||
if line.strip() and not line.lstrip().startswith("#")
|
||||
}
|
||||
|
||||
assert {".git", ".venv", "datasets", "models", "runs", "mlruns"} <= ignored
|
||||
271
tests/test_export_runner.py
Normal file
271
tests/test_export_runner.py
Normal file
|
|
@ -0,0 +1,271 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import ModuleType
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from yolo_webui import export_runner
|
||||
|
||||
|
||||
def _write_config(tmp_path: Path, data: Any) -> Path:
|
||||
path = tmp_path / "export.json"
|
||||
path.write_text(json.dumps(data), encoding="utf-8")
|
||||
return path
|
||||
|
||||
|
||||
def _write_model(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
models = tmp_path / "models"
|
||||
models.mkdir()
|
||||
model = models / "model.pt"
|
||||
model.write_bytes(b"checkpoint")
|
||||
return model
|
||||
|
||||
|
||||
def _install_fake_ultralytics(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
result: Any,
|
||||
) -> tuple[list[str], list[dict[str, Any]]]:
|
||||
constructed: list[str] = []
|
||||
export_calls: list[dict[str, Any]] = []
|
||||
|
||||
class FakeYOLO:
|
||||
def __init__(self, model: str) -> None:
|
||||
assert os.environ["ULTRALYTICS_SAFE_LOAD"] == "1"
|
||||
constructed.append(model)
|
||||
|
||||
def export(self, **kwargs: Any) -> Any:
|
||||
export_calls.append(kwargs)
|
||||
return result
|
||||
|
||||
fake_ultralytics = ModuleType("ultralytics")
|
||||
fake_ultralytics.YOLO = FakeYOLO # type: ignore[attr-defined]
|
||||
monkeypatch.setitem(sys.modules, "ultralytics", fake_ultralytics)
|
||||
return constructed, export_calls
|
||||
|
||||
|
||||
def test_main_validates_config_and_reports_existing_export(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
model = _write_model(tmp_path, monkeypatch)
|
||||
result = tmp_path / "models" / "model.onnx"
|
||||
result.write_bytes(b"onnx")
|
||||
constructed, export_calls = _install_fake_ultralytics(monkeypatch, result)
|
||||
monkeypatch.setenv("ULTRALYTICS_SAFE_LOAD", "0")
|
||||
config = _write_config(
|
||||
tmp_path,
|
||||
{
|
||||
"model": model.name,
|
||||
"format": "ONNX",
|
||||
"imgsz": 320,
|
||||
"half": False,
|
||||
"int8": False,
|
||||
"dynamic": True,
|
||||
"simplify": True,
|
||||
"batch": 2,
|
||||
"workspace": 2.5,
|
||||
},
|
||||
)
|
||||
|
||||
return_code = export_runner.main([str(config)])
|
||||
|
||||
output = capsys.readouterr()
|
||||
assert return_code == 0
|
||||
assert constructed == [str(model)]
|
||||
assert export_calls == [
|
||||
{
|
||||
"format": "onnx",
|
||||
"imgsz": 320,
|
||||
"half": False,
|
||||
"int8": False,
|
||||
"dynamic": True,
|
||||
"simplify": True,
|
||||
"batch": 2,
|
||||
}
|
||||
]
|
||||
assert f"__YOLO_WEBUI_RESULT__:{result}" in output.out
|
||||
|
||||
|
||||
def test_main_passes_workspace_only_for_engine_format(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
model = _write_model(tmp_path, monkeypatch)
|
||||
result = tmp_path / "models" / "model.engine"
|
||||
result.write_bytes(b"engine")
|
||||
_, export_calls = _install_fake_ultralytics(monkeypatch, result)
|
||||
config = _write_config(
|
||||
tmp_path,
|
||||
{
|
||||
"model": model.name,
|
||||
"format": "engine",
|
||||
"imgsz": 640,
|
||||
"half": False,
|
||||
"int8": False,
|
||||
"dynamic": False,
|
||||
"simplify": True,
|
||||
"batch": 1,
|
||||
"workspace": 4.0,
|
||||
},
|
||||
)
|
||||
|
||||
return_code = export_runner.main([str(config)])
|
||||
assert return_code == 0
|
||||
assert export_calls[0]["workspace"] == 4.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"override",
|
||||
[
|
||||
{"format": "onnxx"},
|
||||
{"format": ["onnx"]},
|
||||
{"imgsz": "640"},
|
||||
{"imgsz": True},
|
||||
{"half": 1},
|
||||
{"batch": 0},
|
||||
{"workspace": float("inf")},
|
||||
{"workspace": 10**400},
|
||||
{"workspace": 64.5},
|
||||
{"unexpected": "value"},
|
||||
],
|
||||
)
|
||||
def test_invalid_export_values_are_rejected_before_loading_ultralytics(
|
||||
override: dict[str, Any],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
model = _write_model(tmp_path, monkeypatch)
|
||||
constructed, _ = _install_fake_ultralytics(monkeypatch, model)
|
||||
config = _write_config(
|
||||
tmp_path,
|
||||
{
|
||||
"model": str(model),
|
||||
"format": "onnx",
|
||||
**override,
|
||||
},
|
||||
)
|
||||
|
||||
return_code = export_runner.main([str(config)])
|
||||
|
||||
output = capsys.readouterr()
|
||||
assert return_code == 1
|
||||
assert constructed == []
|
||||
assert "Ошибка конфигурации экспорта:" in output.err
|
||||
|
||||
|
||||
def test_non_object_json_is_reported_as_invalid_config(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
|
||||
return_code = export_runner.main([str(_write_config(tmp_path, ["model.pt"]))])
|
||||
|
||||
output = capsys.readouterr()
|
||||
assert return_code == 1
|
||||
assert "JSON-объектом" in output.err
|
||||
|
||||
|
||||
def test_model_must_be_a_regular_file_inside_an_allowed_root(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
(tmp_path / "models").mkdir()
|
||||
outside = tmp_path / "outside.pt"
|
||||
outside.write_bytes(b"checkpoint")
|
||||
|
||||
return_code = export_runner.main(
|
||||
[str(_write_config(tmp_path, {"model": str(outside), "format": "onnx"}))]
|
||||
)
|
||||
|
||||
output = capsys.readouterr()
|
||||
assert return_code == 1
|
||||
assert "разрешённом каталоге" in output.err
|
||||
|
||||
|
||||
def test_symlink_cannot_escape_an_allowed_model_root(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
models = tmp_path / "models"
|
||||
models.mkdir()
|
||||
outside = tmp_path / "outside.pt"
|
||||
outside.write_bytes(b"checkpoint")
|
||||
link = models / "linked.pt"
|
||||
link.symlink_to(outside)
|
||||
|
||||
return_code = export_runner.main(
|
||||
[str(_write_config(tmp_path, {"model": str(link), "format": "onnx"}))]
|
||||
)
|
||||
|
||||
output = capsys.readouterr()
|
||||
assert return_code == 1
|
||||
assert "разрешённом каталоге" in output.err
|
||||
|
||||
|
||||
def test_half_and_int8_cannot_be_enabled_together(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
model = _write_model(tmp_path, monkeypatch)
|
||||
|
||||
return_code = export_runner.main(
|
||||
[
|
||||
str(
|
||||
_write_config(
|
||||
tmp_path,
|
||||
{
|
||||
"model": str(model),
|
||||
"format": "onnx",
|
||||
"half": True,
|
||||
"int8": True,
|
||||
},
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
output = capsys.readouterr()
|
||||
assert return_code == 1
|
||||
assert "нельзя включать одновременно" in output.err
|
||||
|
||||
|
||||
@pytest.mark.parametrize("result", [None, [], "/missing/export.onnx"])
|
||||
def test_missing_export_artifact_is_a_failure(
|
||||
result: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
model = _write_model(tmp_path, monkeypatch)
|
||||
_install_fake_ultralytics(monkeypatch, result)
|
||||
|
||||
return_code = export_runner.main(
|
||||
[
|
||||
str(
|
||||
_write_config(
|
||||
tmp_path,
|
||||
{"model": str(model), "format": "onnx"},
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
output = capsys.readouterr()
|
||||
assert return_code == 1
|
||||
assert "__YOLO_WEBUI_RESULT__:" not in output.out
|
||||
assert "Ошибка при экспорте модели:" in output.out
|
||||
|
|
@ -1,5 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from html.parser import HTMLParser
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
import subprocess
|
||||
|
|
@ -9,6 +10,37 @@ import pytest
|
|||
|
||||
NODE = shutil.which("node")
|
||||
APP_JS = Path("src/yolo_webui/static/app.js")
|
||||
INDEX_HTML = Path("src/yolo_webui/static/index.html")
|
||||
STYLE_CSS = Path("src/yolo_webui/static/style.css")
|
||||
|
||||
|
||||
class _MarkupCollector(HTMLParser):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.by_id: dict[str, dict[str, str | None]] = {}
|
||||
self.options: dict[str, str] = {}
|
||||
self._option_value: str | None = None
|
||||
|
||||
def handle_starttag(
|
||||
self,
|
||||
tag: str,
|
||||
attrs: list[tuple[str, str | None]],
|
||||
) -> None:
|
||||
attributes = dict(attrs)
|
||||
if element_id := attributes.get("id"):
|
||||
self.by_id[element_id] = attributes
|
||||
if tag == "option":
|
||||
self._option_value = attributes.get("value")
|
||||
|
||||
def handle_data(self, data: str) -> None:
|
||||
if self._option_value is not None:
|
||||
self.options[self._option_value] = (
|
||||
self.options.get(self._option_value, "") + data
|
||||
).strip()
|
||||
|
||||
def handle_endtag(self, tag: str) -> None:
|
||||
if tag == "option":
|
||||
self._option_value = None
|
||||
|
||||
|
||||
@pytest.mark.skipif(NODE is None, reason="Node.js is required for frontend checks")
|
||||
|
|
@ -17,5 +49,30 @@ def test_frontend_javascript_syntax() -> None:
|
|||
|
||||
|
||||
@pytest.mark.skipif(NODE is None, reason="Node.js is required for frontend checks")
|
||||
def test_frontend_zero_values_and_dynamic_chart_series() -> None:
|
||||
def test_frontend_interactions_smoke() -> None:
|
||||
subprocess.run([NODE, "tests/frontend_smoke.js"], check=True)
|
||||
|
||||
|
||||
def test_export_controls_have_correct_formats_constraints_and_styles() -> None:
|
||||
parser = _MarkupCollector()
|
||||
parser.feed(INDEX_HTML.read_text(encoding="utf-8"))
|
||||
|
||||
assert parser.options["saved_model"].startswith("TensorFlow SavedModel")
|
||||
assert parser.options["pb"] == "TensorFlow GraphDef (.pb)"
|
||||
assert "triton" not in parser.options
|
||||
assert parser.by_id["export-imgsz"]["min"] == "32"
|
||||
assert parser.by_id["export-imgsz"]["max"] == "8192"
|
||||
assert parser.by_id["export-imgsz"]["step"] == "1"
|
||||
assert parser.by_id["export-batch"]["min"] == "1"
|
||||
assert parser.by_id["export-batch"]["max"] == "1024"
|
||||
assert parser.by_id["export-batch"]["step"] == "1"
|
||||
assert parser.by_id["export-workspace"]["min"] == "1"
|
||||
assert parser.by_id["export-workspace"]["max"] == "64"
|
||||
assert parser.by_id["export-workspace"]["step"] == "0.5"
|
||||
|
||||
css = STYLE_CSS.read_text(encoding="utf-8")
|
||||
assert "#export-config-pane" in css
|
||||
assert "#export-run-pane" in css
|
||||
assert "#export-log-container::-webkit-scrollbar" in css
|
||||
assert "#export-status-card.status-exporting" in css
|
||||
assert "#export-status-card.status-failed" in css
|
||||
|
|
|
|||
423
tests/test_mlflow_metrics.py
Normal file
423
tests/test_mlflow_metrics.py
Normal file
|
|
@ -0,0 +1,423 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import math
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from yolo_webui.config import YoloTask
|
||||
from yolo_webui.mlflow_metrics import (
|
||||
PER_CLASS_FIELDS,
|
||||
TaskMetricsMonitor,
|
||||
collect_monitoring_metrics,
|
||||
collect_per_class_metrics,
|
||||
)
|
||||
|
||||
|
||||
def _component(offset: float = 0.0) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
mp=0.75 + offset,
|
||||
mr=0.65 + offset,
|
||||
p=[0.8 + offset, 0.7 + offset],
|
||||
r=[0.6 + offset, 0.7 + offset],
|
||||
f1=[0.685714 + offset, 0.7 + offset],
|
||||
map50=0.72 + offset,
|
||||
map75=0.58 + offset,
|
||||
map=0.61 + offset,
|
||||
ap50=[0.76 + offset, 0.68 + offset],
|
||||
ap=[0.64 + offset, 0.58 + offset],
|
||||
ap_class_index=[0, 1],
|
||||
)
|
||||
|
||||
|
||||
def _metric_set(task: YoloTask) -> SimpleNamespace:
|
||||
common: dict[str, Any] = {
|
||||
"fitness": 0.61,
|
||||
"speed": {
|
||||
"preprocess": 0.1,
|
||||
"inference": 1.2,
|
||||
"loss": 0.3,
|
||||
"postprocess": 0.4,
|
||||
},
|
||||
"names": {0: "cat", 1: "dog"},
|
||||
"nt_per_class": [7, 5],
|
||||
}
|
||||
if task == "classify":
|
||||
return SimpleNamespace(
|
||||
**common,
|
||||
top1=0.8,
|
||||
top5=0.95,
|
||||
confusion_matrix=SimpleNamespace(
|
||||
matrix=[
|
||||
[7, 2],
|
||||
[1, 4],
|
||||
],
|
||||
names=common["names"],
|
||||
),
|
||||
)
|
||||
|
||||
common["box"] = _component()
|
||||
if task == "segment":
|
||||
common["seg"] = _component(0.05)
|
||||
elif task == "pose":
|
||||
common["pose"] = _component(0.1)
|
||||
return SimpleNamespace(**common)
|
||||
|
||||
|
||||
def _trainer(task: YoloTask, save_dir: Path | None = None) -> SimpleNamespace:
|
||||
def label_loss_items(
|
||||
loss: object,
|
||||
prefix: str = "train",
|
||||
) -> dict[str, float] | list[str]:
|
||||
keys = [f"{prefix}/first_loss", f"{prefix}/second_loss"]
|
||||
return keys if loss is None else dict(zip(keys, (0.25, 0.5)))
|
||||
|
||||
return SimpleNamespace(
|
||||
metrics={
|
||||
"val/first_loss": 0.2,
|
||||
"val/second_loss": 0.3,
|
||||
"val/future_non_loss_metric": 99.0,
|
||||
},
|
||||
validator=SimpleNamespace(metrics=_metric_set(task), training=True),
|
||||
fitness=0.6,
|
||||
best_fitness=0.65,
|
||||
tloss=[0.25, 0.5],
|
||||
label_loss_items=label_loss_items,
|
||||
lr={"lr/pg0": 0.01, "lr/pg1": 0.02},
|
||||
epoch_time=12.5,
|
||||
save_dir=save_dir,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("task", "required"),
|
||||
[
|
||||
(
|
||||
"detect",
|
||||
{
|
||||
"monitor/quality/box/f1_at_optimal_confidence",
|
||||
"monitor/quality/box/map75",
|
||||
"monitor/quality/box/worst_class_map50_95",
|
||||
},
|
||||
),
|
||||
(
|
||||
"segment",
|
||||
{
|
||||
"monitor/quality/box/f1_at_optimal_confidence",
|
||||
"monitor/quality/mask/f1_at_optimal_confidence",
|
||||
"monitor/quality/mask/map75",
|
||||
},
|
||||
),
|
||||
(
|
||||
"pose",
|
||||
{
|
||||
"monitor/quality/box/f1_at_optimal_confidence",
|
||||
"monitor/quality/keypoints/f1_at_optimal_confidence",
|
||||
"monitor/quality/keypoints/map75",
|
||||
},
|
||||
),
|
||||
(
|
||||
"obb",
|
||||
{
|
||||
"monitor/quality/oriented_box/f1_at_optimal_confidence",
|
||||
"monitor/quality/oriented_box/map75",
|
||||
"monitor/quality/oriented_box/worst_class_map50_95",
|
||||
},
|
||||
),
|
||||
(
|
||||
"classify",
|
||||
{
|
||||
"monitor/quality/classification/top1_error",
|
||||
"monitor/quality/classification/macro_f1",
|
||||
"monitor/quality/classification/balanced_accuracy",
|
||||
"monitor/quality/classification/worst_class_recall",
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_collect_monitoring_metrics_has_task_specific_contract(
|
||||
task: YoloTask,
|
||||
required: set[str],
|
||||
) -> None:
|
||||
metrics = collect_monitoring_metrics(task, _trainer(task))
|
||||
|
||||
assert required <= metrics.keys()
|
||||
assert metrics["monitor/fitness/ultralytics_current"] == pytest.approx(0.61)
|
||||
assert metrics["monitor/fitness/ultralytics_best"] == pytest.approx(0.65)
|
||||
expected_scale = 2 if task in {"segment", "pose"} else 1
|
||||
assert metrics["monitor/fitness/task_score_current"] == pytest.approx(
|
||||
0.61 / expected_scale
|
||||
)
|
||||
assert metrics["monitor/fitness/task_score_best"] == pytest.approx(
|
||||
0.65 / expected_scale
|
||||
)
|
||||
assert metrics["monitor/loss/train_total"] == pytest.approx(0.75)
|
||||
assert metrics["monitor/loss/validation_total"] == pytest.approx(0.5)
|
||||
assert metrics["monitor/loss/generalization_gap"] == pytest.approx(-0.25)
|
||||
assert metrics["monitor/loss/validation_to_train_ratio"] == pytest.approx(2 / 3)
|
||||
assert metrics["monitor/optimization/learning_rate_mean"] == pytest.approx(0.015)
|
||||
assert metrics["monitor/performance/epoch_seconds"] == pytest.approx(12.5)
|
||||
assert metrics["monitor/performance/validation_inference_ms_per_image"] == pytest.approx(1.2)
|
||||
|
||||
|
||||
def test_classification_metrics_are_derived_from_confusion_matrix() -> None:
|
||||
metrics = collect_monitoring_metrics("classify", _trainer("classify"))
|
||||
|
||||
assert metrics["monitor/quality/classification/top1_error"] == pytest.approx(0.2)
|
||||
assert metrics["monitor/quality/classification/macro_precision"] == pytest.approx(
|
||||
(7 / 9 + 4 / 5) / 2
|
||||
)
|
||||
assert metrics["monitor/quality/classification/macro_recall"] == pytest.approx(
|
||||
(7 / 8 + 4 / 6) / 2
|
||||
)
|
||||
assert metrics["monitor/quality/classification/classes_evaluated"] == 2
|
||||
|
||||
|
||||
def test_mask_and_keypoint_metrics_use_their_own_components() -> None:
|
||||
segment_metrics = collect_monitoring_metrics("segment", _trainer("segment"))
|
||||
pose_metrics = collect_monitoring_metrics("pose", _trainer("pose"))
|
||||
|
||||
assert segment_metrics[
|
||||
"monitor/quality/mask/f1_at_optimal_confidence"
|
||||
] == pytest.approx((0.735714 + 0.75) / 2)
|
||||
assert pose_metrics[
|
||||
"monitor/quality/keypoints/f1_at_optimal_confidence"
|
||||
] == pytest.approx((0.785714 + 0.8) / 2)
|
||||
assert segment_metrics[
|
||||
"monitor/quality/mask/f1_at_optimal_confidence"
|
||||
] != segment_metrics["monitor/quality/box/f1_at_optimal_confidence"]
|
||||
|
||||
|
||||
def test_classification_macro_metrics_include_false_positive_only_class() -> None:
|
||||
trainer = _trainer("classify")
|
||||
trainer.validator.metrics.confusion_matrix = SimpleNamespace(
|
||||
matrix=[
|
||||
[7, 2, 0],
|
||||
[1, 4, 0],
|
||||
[1, 0, 0],
|
||||
],
|
||||
names={0: "cat", 1: "dog", 2: "fox"},
|
||||
)
|
||||
|
||||
metrics = collect_monitoring_metrics("classify", trainer)
|
||||
rows = collect_per_class_metrics("classify", trainer)
|
||||
|
||||
assert metrics["monitor/quality/classification/classes_evaluated"] == 3
|
||||
assert metrics["monitor/quality/classification/macro_precision"] == pytest.approx(
|
||||
(7 / 9 + 4 / 5 + 0.0) / 3
|
||||
)
|
||||
assert metrics["monitor/quality/classification/macro_recall"] == pytest.approx(
|
||||
(7 / 9 + 4 / 6) / 2
|
||||
)
|
||||
assert rows[-1]["class_name"] == "fox"
|
||||
assert rows[-1]["support"] == 0
|
||||
assert rows[-1]["precision"] == 0
|
||||
|
||||
|
||||
def test_monitor_drops_non_finite_and_non_numeric_values() -> None:
|
||||
trainer = _trainer("detect")
|
||||
trainer.best_fitness = float("nan")
|
||||
trainer.epoch_time = float("inf")
|
||||
trainer.lr = {"lr/pg0": "not-a-number"}
|
||||
trainer.validator.metrics.box.map75 = None
|
||||
|
||||
metrics = collect_monitoring_metrics("detect", trainer)
|
||||
|
||||
assert "monitor/fitness/ultralytics_best" not in metrics
|
||||
assert "monitor/fitness/task_score_best" not in metrics
|
||||
assert "monitor/performance/epoch_seconds" not in metrics
|
||||
assert "monitor/optimization/learning_rate_mean" not in metrics
|
||||
assert "monitor/quality/box/map75" not in metrics
|
||||
assert all(
|
||||
value == value and value not in (float("inf"), float("-inf"))
|
||||
for value in metrics.values()
|
||||
)
|
||||
|
||||
|
||||
def test_callback_enriches_trainer_metrics_before_mlflow_logging() -> None:
|
||||
trainer = _trainer("obb")
|
||||
monitor = TaskMetricsMonitor("obb", mlflow_enabled=True)
|
||||
|
||||
monitor.on_fit_epoch_end(trainer)
|
||||
|
||||
assert trainer.metrics["monitor/fitness/ultralytics_current"] == pytest.approx(0.61)
|
||||
assert trainer.metrics[
|
||||
"monitor/quality/oriented_box/f1_at_optimal_confidence"
|
||||
] == pytest.approx((0.685714 + 0.7) / 2)
|
||||
|
||||
|
||||
def test_final_validation_does_not_repeat_stale_train_diagnostics() -> None:
|
||||
trainer = _trainer("detect")
|
||||
trainer.validator.training = False
|
||||
|
||||
metrics = collect_monitoring_metrics("detect", trainer)
|
||||
|
||||
assert "monitor/quality/box/f1_at_optimal_confidence" in metrics
|
||||
assert "monitor/performance/validation_inference_ms_per_image" in metrics
|
||||
assert "monitor/loss/train_total" not in metrics
|
||||
assert "monitor/loss/validation_total" not in metrics
|
||||
assert "monitor/optimization/learning_rate_mean" not in metrics
|
||||
assert "monitor/performance/epoch_seconds" not in metrics
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("task", "expected_components"),
|
||||
[
|
||||
("detect", {"box"}),
|
||||
("segment", {"box", "mask"}),
|
||||
("pose", {"box", "keypoints"}),
|
||||
("obb", {"oriented_box"}),
|
||||
("classify", {"classification"}),
|
||||
],
|
||||
)
|
||||
def test_final_per_class_rows_cover_every_task_component(
|
||||
task: YoloTask,
|
||||
expected_components: set[str],
|
||||
) -> None:
|
||||
rows = collect_per_class_metrics(task, _trainer(task))
|
||||
|
||||
assert {row["component"] for row in rows} == expected_components
|
||||
assert {row["class_name"] for row in rows} == {"cat", "dog"}
|
||||
assert all(row["support"] > 0 for row in rows)
|
||||
|
||||
|
||||
def test_per_class_rows_preserve_positions_of_non_finite_values() -> None:
|
||||
trainer = _trainer("detect")
|
||||
component = trainer.validator.metrics.box
|
||||
component.ap_class_index = [0, 1, 2]
|
||||
component.p = [0.8, float("nan"), 0.6]
|
||||
component.r = [0.7, 0.5, 0.4]
|
||||
component.f1 = [0.746, 0.5, 0.48]
|
||||
component.ap50 = [0.75, 0.55, 0.45]
|
||||
component.ap = [0.65, 0.45, 0.35]
|
||||
trainer.validator.metrics.names = {0: "first", 1: "second", 2: "third"}
|
||||
trainer.validator.metrics.nt_per_class = [5, 4, 3]
|
||||
|
||||
rows = collect_per_class_metrics("detect", trainer)
|
||||
|
||||
assert [row["class_id"] for row in rows] == [0, 1, 2]
|
||||
assert rows[1]["precision"] == ""
|
||||
assert rows[2]["precision"] == pytest.approx(0.6)
|
||||
|
||||
|
||||
def test_train_end_writes_mlflow_uploadable_per_class_csv(tmp_path: Path) -> None:
|
||||
trainer = _trainer("segment", tmp_path)
|
||||
monitor = TaskMetricsMonitor("segment", mlflow_enabled=True)
|
||||
|
||||
monitor.on_train_end(trainer)
|
||||
|
||||
output = tmp_path / "monitoring" / "task_metrics.csv"
|
||||
assert output.exists()
|
||||
with output.open(encoding="utf-8") as stream:
|
||||
reader = csv.DictReader(stream)
|
||||
rows = list(reader)
|
||||
assert tuple(reader.fieldnames or ()) == PER_CLASS_FIELDS
|
||||
assert len(rows) == 4
|
||||
assert {row["component"] for row in rows} == {"box", "mask"}
|
||||
assert {row["task"] for row in rows} == {"segment"}
|
||||
assert {row["class_id"] for row in rows} == {"0", "1"}
|
||||
assert {row["class_name"] for row in rows} == {"cat", "dog"}
|
||||
for row in rows:
|
||||
for field in ("support", "precision", "recall", "f1", "map50", "map50_95"):
|
||||
assert math.isfinite(float(row[field]))
|
||||
|
||||
|
||||
def test_disabled_mlflow_does_not_add_metrics_or_write_artifact(tmp_path: Path) -> None:
|
||||
trainer = _trainer("detect", tmp_path)
|
||||
original_metrics = dict(trainer.metrics)
|
||||
monitor = TaskMetricsMonitor("detect", mlflow_enabled=False)
|
||||
|
||||
monitor.on_fit_epoch_end(trainer)
|
||||
monitor.on_train_end(trainer)
|
||||
|
||||
assert trainer.metrics == original_metrics
|
||||
assert not (tmp_path / "monitoring" / "task_metrics.csv").exists()
|
||||
|
||||
|
||||
def test_train_end_writes_header_when_no_classes_were_evaluated(tmp_path: Path) -> None:
|
||||
trainer = _trainer("detect", tmp_path)
|
||||
component = trainer.validator.metrics.box
|
||||
for attribute in ("p", "r", "f1", "ap50", "ap", "ap_class_index"):
|
||||
setattr(component, attribute, [])
|
||||
|
||||
TaskMetricsMonitor("detect", mlflow_enabled=True).on_train_end(trainer)
|
||||
|
||||
output = tmp_path / "monitoring" / "task_metrics.csv"
|
||||
with output.open(encoding="utf-8") as stream:
|
||||
reader = csv.DictReader(stream)
|
||||
rows = list(reader)
|
||||
assert tuple(reader.fieldnames or ()) == PER_CLASS_FIELDS
|
||||
assert rows == []
|
||||
|
||||
|
||||
def test_monitoring_contract_is_persisted_by_mlflow(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import mlflow
|
||||
from mlflow.tracking import MlflowClient
|
||||
from ultralytics.utils.callbacks import mlflow as ultralytics_mlflow
|
||||
from yolo_webui.ultralytics_trainers import trainer_for_task
|
||||
|
||||
previous_uri = mlflow.get_tracking_uri()
|
||||
tracking_uri = f"sqlite:///{tmp_path / 'mlflow.db'}"
|
||||
run_id = ""
|
||||
try:
|
||||
monkeypatch.setenv("YOLO_WEBUI_MLFLOW_RUN_GROUP", "test-group")
|
||||
mlflow.set_tracking_uri(tracking_uri)
|
||||
client = MlflowClient(tracking_uri=tracking_uri)
|
||||
experiment_id = client.create_experiment(
|
||||
"task-monitor-test",
|
||||
artifact_location=(tmp_path / "artifacts").as_uri(),
|
||||
)
|
||||
trainer_class = trainer_for_task("obb")
|
||||
trainer = trainer_class.__new__(trainer_class)
|
||||
trainer.__dict__.update(vars(_trainer("obb", tmp_path / "run")))
|
||||
trainer._mlflow_active = True
|
||||
trainer.epoch = 0
|
||||
trainer.callbacks = {
|
||||
"on_train_start": [],
|
||||
"on_fit_epoch_end": [ultralytics_mlflow.on_fit_epoch_end],
|
||||
"on_train_end": [],
|
||||
}
|
||||
trainer._install_task_metrics_monitor()
|
||||
trainer._task_metrics_monitor.mlflow_enabled = True
|
||||
|
||||
with mlflow.start_run(experiment_id=experiment_id, run_name="obb-contract") as run:
|
||||
run_id = run.info.run_id
|
||||
for callback in trainer.callbacks["on_train_start"]:
|
||||
callback(trainer)
|
||||
for callback in trainer.callbacks["on_fit_epoch_end"]:
|
||||
callback(trainer)
|
||||
|
||||
trainer.epoch = 1
|
||||
trainer.validator.training = False
|
||||
trainer.metrics = {"metrics/mAP50": 0.72}
|
||||
for callback in trainer.callbacks["on_fit_epoch_end"]:
|
||||
callback(trainer)
|
||||
for callback in trainer.callbacks["on_train_end"]:
|
||||
callback(trainer)
|
||||
|
||||
persisted = client.get_run(run_id)
|
||||
history = client.get_metric_history(
|
||||
run_id,
|
||||
"monitor/quality/oriented_box/f1_at_optimal_confidence",
|
||||
)
|
||||
assert persisted.data.tags["yolo.task"] == "obb"
|
||||
assert persisted.data.tags["monitoring.schema_version"] == "1"
|
||||
assert persisted.data.tags["yolo.run_group"] == "test-group"
|
||||
assert persisted.data.metrics[
|
||||
"monitor/fitness/ultralytics_current"
|
||||
] == pytest.approx(0.61)
|
||||
assert [point.step for point in history] == [0, 1]
|
||||
assert {
|
||||
artifact.path
|
||||
for artifact in client.list_artifacts(run_id, "monitoring")
|
||||
} == {"monitoring/task_metrics.csv"}
|
||||
finally:
|
||||
if mlflow.active_run() is not None:
|
||||
mlflow.end_run(status="KILLED")
|
||||
mlflow.set_tracking_uri(previous_uri)
|
||||
43
tests/test_mlflow_smoke_verifier.py
Normal file
43
tests/test_mlflow_smoke_verifier.py
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
SPEC = importlib.util.spec_from_file_location(
|
||||
"verify_mlflow_smoke",
|
||||
Path(__file__).parents[1] / "scripts" / "verify_mlflow_smoke.py",
|
||||
)
|
||||
assert SPEC is not None and SPEC.loader is not None
|
||||
verify_mlflow_smoke = importlib.util.module_from_spec(SPEC)
|
||||
SPEC.loader.exec_module(verify_mlflow_smoke)
|
||||
latest_smoke_run_group = verify_mlflow_smoke.latest_smoke_run_group
|
||||
latest_task_run = verify_mlflow_smoke.latest_task_run
|
||||
|
||||
|
||||
def _run(name: str, group: str, *, anchor: bool = False) -> SimpleNamespace:
|
||||
tags = {
|
||||
"mlflow.runName": name,
|
||||
"yolo.run_group": group,
|
||||
}
|
||||
if anchor:
|
||||
tags["smoke.anchor"] = "true"
|
||||
return SimpleNamespace(data=SimpleNamespace(tags=tags))
|
||||
|
||||
|
||||
def test_verifier_selects_one_complete_smoke_batch() -> None:
|
||||
runs = [
|
||||
_run("detect-smoke", "new"),
|
||||
_run("smoke-batch-new", "new", anchor=True),
|
||||
_run("obb-smoke", "old"),
|
||||
_run("smoke-batch-old", "old", anchor=True),
|
||||
]
|
||||
|
||||
group = latest_smoke_run_group(runs)
|
||||
|
||||
assert group == "new"
|
||||
assert latest_task_run(runs, "detect", group).data.tags["yolo.run_group"] == "new"
|
||||
with pytest.raises(AssertionError, match="obb-smoke"):
|
||||
latest_task_run(runs, "obb", group)
|
||||
|
|
@ -35,6 +35,14 @@ def test_read_classes_custom_yaml(
|
|||
assert read_classes(tmp_path, str(custom_file)) == expected
|
||||
|
||||
|
||||
def test_read_classes_rejects_fractional_yaml_class_id(tmp_path: Path) -> None:
|
||||
custom_file = tmp_path / "custom_classes.yaml"
|
||||
custom_file.write_text("names:\n 0.5: person\n", encoding="utf-8")
|
||||
|
||||
with pytest.raises(ValueError, match="Некорректный ID класса"):
|
||||
read_classes(tmp_path, str(custom_file))
|
||||
|
||||
|
||||
def test_missing_custom_classes_path_does_not_fall_back(tmp_path: Path) -> None:
|
||||
(tmp_path / "classes.txt").write_text("fallback\n", encoding="utf-8")
|
||||
|
||||
|
|
@ -145,6 +153,59 @@ def test_split_dataset_rejects_single_image_without_writing_output(
|
|||
assert not (tmp_path / ".yolo-webui").exists()
|
||||
|
||||
|
||||
def test_split_dataset_rejects_label_id_missing_from_classes(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
images_dir = tmp_path / "images"
|
||||
labels_dir = tmp_path / "labels"
|
||||
images_dir.mkdir()
|
||||
labels_dir.mkdir()
|
||||
for index in range(2):
|
||||
(images_dir / f"image-{index}.jpg").write_bytes(b"")
|
||||
(labels_dir / f"image-{index}.txt").write_text(
|
||||
"1 0.5 0.5 0.2 0.2\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
(tmp_path / "classes.txt").write_text("only-class\n", encoding="utf-8")
|
||||
|
||||
with pytest.raises(ValueError, match="указан класс 1"):
|
||||
split_dataset(str(tmp_path), 0.5, "")
|
||||
|
||||
assert not (tmp_path / ".yolo-webui").exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("invalid_id", ["-1", "0.5", "nan", "class-a"])
|
||||
def test_split_dataset_rejects_invalid_label_class_id(
|
||||
tmp_path: Path,
|
||||
invalid_id: str,
|
||||
) -> None:
|
||||
images_dir = tmp_path / "images"
|
||||
labels_dir = tmp_path / "labels"
|
||||
images_dir.mkdir()
|
||||
labels_dir.mkdir()
|
||||
for index in range(2):
|
||||
(images_dir / f"image-{index}.jpg").write_bytes(b"")
|
||||
(labels_dir / f"image-{index}.txt").write_text(
|
||||
f"{invalid_id} 0.5 0.5 0.2 0.2\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
(tmp_path / "classes.txt").write_text("item\n", encoding="utf-8")
|
||||
|
||||
with pytest.raises(ValueError, match="Некорректный ID класса"):
|
||||
split_dataset(str(tmp_path), 0.5, "")
|
||||
|
||||
assert not (tmp_path / ".yolo-webui").exists()
|
||||
|
||||
|
||||
def test_split_dataset_rejects_non_finite_ratio_before_writing(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
with pytest.raises(ValueError, match="конечным числом"):
|
||||
split_dataset(str(tmp_path), float("nan"), "")
|
||||
|
||||
assert not (tmp_path / ".yolo-webui").exists()
|
||||
|
||||
|
||||
def test_split_dataset_finds_nested_images_and_labels(tmp_path: Path) -> None:
|
||||
images_dir = tmp_path / "images" / "day"
|
||||
labels_dir = tmp_path / "labels" / "day"
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ from pathlib import Path
|
|||
from types import ModuleType, SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from yolo_webui.config import MlflowConfig, TrainingConfig
|
||||
from yolo_webui.trainer import TrainingEvent, TrainingRunner
|
||||
|
||||
|
|
@ -37,37 +39,59 @@ def test_runner_wires_yolo_callbacks_and_returns_output(
|
|||
train_arguments: list[dict[str, Any]] = []
|
||||
|
||||
class FakeSettings:
|
||||
def __init__(self) -> None:
|
||||
self.values = {"mlflow": True}
|
||||
|
||||
def __getitem__(self, key: str) -> bool:
|
||||
return self.values[key]
|
||||
|
||||
def update(self, values: dict[str, bool]) -> None:
|
||||
settings_updates.append(values)
|
||||
self.values.update(values)
|
||||
|
||||
class FakeYOLO:
|
||||
def __init__(self, model: str, task: str) -> None:
|
||||
constructed.append((model, task))
|
||||
self.callbacks: dict[str, Any] = {}
|
||||
self.callbacks: dict[str, list[Any]] = {}
|
||||
self.trainer = SimpleNamespace(
|
||||
args=SimpleNamespace(epochs=2),
|
||||
epoch=0,
|
||||
metrics={},
|
||||
fitness=0.0,
|
||||
best_fitness=0.0,
|
||||
stop=False,
|
||||
save_dir=tmp_path / "run",
|
||||
)
|
||||
|
||||
def add_callback(self, name: str, callback: Any) -> None:
|
||||
self.callbacks[name] = callback
|
||||
self.callbacks.setdefault(name, []).append(callback)
|
||||
|
||||
def run_callbacks(self, name: str) -> None:
|
||||
for callback in self.callbacks.get(name, []):
|
||||
callback(self.trainer)
|
||||
|
||||
def train(self, **kwargs: Any) -> None:
|
||||
train_arguments.append(kwargs)
|
||||
self.callbacks["on_train_start"](self.trainer)
|
||||
self.run_callbacks("on_train_start")
|
||||
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)
|
||||
self.trainer.fitness = 0.5 + epoch / 10
|
||||
self.trainer.best_fitness = self.trainer.fitness
|
||||
self.run_callbacks("on_fit_epoch_end")
|
||||
self.run_callbacks("on_train_end")
|
||||
|
||||
fake_ultralytics = ModuleType("ultralytics")
|
||||
fake_ultralytics.YOLO = FakeYOLO # type: ignore[attr-defined]
|
||||
fake_ultralytics.settings = FakeSettings() # type: ignore[attr-defined]
|
||||
fake_trainers = ModuleType("yolo_webui.ultralytics_trainers")
|
||||
|
||||
class FakeTrainer:
|
||||
pass
|
||||
|
||||
fake_trainers.trainer_for_task = lambda _task: FakeTrainer # type: ignore[attr-defined]
|
||||
monkeypatch.setitem(sys.modules, "ultralytics", fake_ultralytics)
|
||||
monkeypatch.setitem(sys.modules, "yolo_webui.ultralytics_trainers", fake_trainers)
|
||||
|
||||
config = TrainingConfig(
|
||||
dataset="dataset.yaml",
|
||||
|
|
@ -82,10 +106,69 @@ def test_runner_wires_yolo_callbacks_and_returns_output(
|
|||
|
||||
assert output == tmp_path / "run"
|
||||
assert constructed == [("models/model.pt", "pose")]
|
||||
assert settings_updates == [{"mlflow": False}]
|
||||
assert [u for u in settings_updates if "mlflow" in u] == [{"mlflow": False}, {"mlflow": True}]
|
||||
assert train_arguments[0]["data"] == "dataset.yaml"
|
||||
assert train_arguments[0]["verbose"] is True
|
||||
assert train_arguments[0]["trainer"] is FakeTrainer
|
||||
assert [event.kind for event in events] == ["info", "started", "epoch", "epoch", "success"]
|
||||
assert "mAP50=0.5" in events[2].message
|
||||
assert "mAP50=0.6" in events[3].message
|
||||
|
||||
|
||||
def test_ultralytics_mlflow_setting_is_restored_when_model_loading_fails(
|
||||
monkeypatch: Any,
|
||||
) -> None:
|
||||
settings_updates: list[bool] = []
|
||||
|
||||
class FakeSettings:
|
||||
def __init__(self) -> None:
|
||||
self.mlflow = True
|
||||
|
||||
def __getitem__(self, key: str) -> bool:
|
||||
assert key == "mlflow"
|
||||
return self.mlflow
|
||||
|
||||
def update(self, values: dict[str, bool]) -> None:
|
||||
self.mlflow = values["mlflow"]
|
||||
settings_updates.append(self.mlflow)
|
||||
|
||||
class FailingYOLO:
|
||||
def __init__(self, _model: str, task: str) -> None:
|
||||
assert task == "detect"
|
||||
raise RuntimeError("model load failed")
|
||||
|
||||
fake_ultralytics = ModuleType("ultralytics")
|
||||
fake_ultralytics.YOLO = FailingYOLO # type: ignore[attr-defined]
|
||||
fake_ultralytics.settings = FakeSettings() # type: ignore[attr-defined]
|
||||
fake_trainers = ModuleType("yolo_webui.ultralytics_trainers")
|
||||
fake_trainers.trainer_for_task = lambda _task: object # type: ignore[attr-defined]
|
||||
monkeypatch.setitem(sys.modules, "ultralytics", fake_ultralytics)
|
||||
monkeypatch.setitem(sys.modules, "yolo_webui.ultralytics_trainers", fake_trainers)
|
||||
config = TrainingConfig(
|
||||
dataset="dataset.yaml",
|
||||
model="model.pt",
|
||||
mlflow=MlflowConfig(enabled=False),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="model load failed"):
|
||||
TrainingRunner().train(config, lambda _event: None)
|
||||
|
||||
assert settings_updates == [False, True]
|
||||
assert fake_ultralytics.settings["mlflow"] is True # type: ignore[index]
|
||||
|
||||
|
||||
def test_final_checkpoint_validation_is_not_reported_as_an_extra_epoch() -> None:
|
||||
events: list[TrainingEvent] = []
|
||||
trainer = SimpleNamespace(
|
||||
args=SimpleNamespace(epochs=1),
|
||||
epoch=1,
|
||||
metrics={"metrics/mAP50": 0.75},
|
||||
validator=SimpleNamespace(training=False),
|
||||
)
|
||||
|
||||
TrainingRunner()._on_epoch_end(events.append)(trainer)
|
||||
|
||||
assert events == []
|
||||
|
||||
|
||||
def test_prepare_run_clears_previous_stop_request() -> None:
|
||||
|
|
|
|||
39
tests/test_ultralytics_trainers.py
Normal file
39
tests/test_ultralytics_trainers.py
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from yolo_webui.config import YoloTask
|
||||
from yolo_webui.mlflow_metrics import TaskMetricsMonitor
|
||||
from yolo_webui.ultralytics_trainers import MONITORED_TRAINERS, trainer_for_task
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task", ["detect", "segment", "classify", "pose", "obb"])
|
||||
def test_every_task_uses_an_importable_monitored_trainer(task: YoloTask) -> None:
|
||||
trainer_class = trainer_for_task(task)
|
||||
|
||||
assert trainer_class is MONITORED_TRAINERS[task]
|
||||
assert trainer_class.monitoring_task == task
|
||||
assert trainer_class.__module__ == "yolo_webui.ultralytics_trainers"
|
||||
|
||||
|
||||
def test_monitor_callbacks_are_prepended_before_integrations() -> None:
|
||||
trainer_class = trainer_for_task("detect")
|
||||
trainer = trainer_class.__new__(trainer_class)
|
||||
integration_callbacks: dict[str, Any] = {
|
||||
"on_train_start": lambda _trainer: None,
|
||||
"on_fit_epoch_end": lambda _trainer: None,
|
||||
"on_train_end": lambda _trainer: None,
|
||||
}
|
||||
trainer.callbacks = {
|
||||
event: [callback]
|
||||
for event, callback in integration_callbacks.items()
|
||||
}
|
||||
|
||||
trainer._install_task_metrics_monitor()
|
||||
|
||||
assert isinstance(trainer._task_metrics_monitor, TaskMetricsMonitor)
|
||||
for event, integration_callback in integration_callbacks.items():
|
||||
assert trainer.callbacks[event][0].__self__ is trainer._task_metrics_monitor
|
||||
assert trainer.callbacks[event][1] is integration_callback
|
||||
Loading…
Add table
Reference in a new issue