533 lines
19 KiB
Python
533 lines
19 KiB
Python
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,
|
||
)
|