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, )