python-basic-pitch/temporal_chord_metrics.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Iterable

from music_token_metrics import (
    eh_sem_acorde,
    normalizar_token,
    qualidade_completa_token,
    qualidade_token,
    raiz_token,
)


@dataclass(frozen=True)
class ChordInterval:
    start: float
    end: float
    label: str

    def __post_init__(self) -> None:
        if self.start < 0 or self.end <= self.start:
            raise ValueError(f"Intervalo invalido: {self.start}..{self.end}")


def validate_timeline(intervals: Iterable[ChordInterval]) -> list[ChordInterval]:
    ordered = sorted(intervals, key=lambda item: (item.start, item.end))
    previous_end = 0.0
    for item in ordered:
        if item.start < previous_end - 1e-9:
            raise ValueError("Intervalos sobrepostos nao sao aceitos")
        previous_end = item.end
    return ordered


def duration_weighted_accuracy(
    reference: Iterable[ChordInterval],
    estimated: Iterable[ChordInterval],
    comparator: Callable[[str, str], bool],
) -> float:
    reference_items = validate_timeline(reference)
    estimated_items = validate_timeline(estimated)
    boundaries = sorted({
        point
        for item in [*reference_items, *estimated_items]
        for point in (item.start, item.end)
    })
    correct = 0.0
    evaluated = 0.0
    for start, end in zip(boundaries, boundaries[1:]):
        if end <= start:
            continue
        midpoint = (start + end) / 2.0
        expected = label_at(reference_items, midpoint)
        predicted = label_at(estimated_items, midpoint)
        duration = end - start
        evaluated += duration
        if comparator(expected, predicted):
            correct += duration
    return correct / evaluated if evaluated else 1.0


def boundary_f1(
    reference: Iterable[ChordInterval],
    estimated: Iterable[ChordInterval],
    tolerance_seconds: float = 0.5,
) -> dict[str, float | int]:
    if tolerance_seconds < 0:
        raise ValueError("A tolerancia deve ser nao negativa")
    expected = internal_boundaries(validate_timeline(reference))
    predicted = internal_boundaries(validate_timeline(estimated))
    unmatched = set(range(len(predicted)))
    matches = 0
    for boundary in expected:
        candidates = [
            index for index in unmatched
            if abs(predicted[index] - boundary) <= tolerance_seconds
        ]
        if not candidates:
            continue
        best = min(candidates, key=lambda index: abs(predicted[index] - boundary))
        unmatched.remove(best)
        matches += 1
    precision = matches / len(predicted) if predicted else (1.0 if not expected else 0.0)
    recall = matches / len(expected) if expected else (1.0 if not predicted else 0.0)
    f1 = (2 * precision * recall / (precision + recall)) if precision + recall else 0.0
    return {
        "matches": matches,
        "reference_boundaries": len(expected),
        "estimated_boundaries": len(predicted),
        "precision": precision,
        "recall": recall,
        "f1": f1,
    }


def evaluate_temporal_chords(
    reference: Iterable[ChordInterval],
    estimated: Iterable[ChordInterval],
    tolerance_seconds: float = 0.5,
) -> dict[str, object]:
    reference_items = validate_timeline(reference)
    estimated_items = validate_timeline(estimated)
    return {
        "exact_wcsr": duration_weighted_accuracy(
            reference_items, estimated_items,
            lambda left, right: normalizar_token(left) == normalizar_token(right),
        ),
        "root_wcsr": duration_weighted_accuracy(
            reference_items, estimated_items,
            lambda left, right: raiz_token(left) == raiz_token(right),
        ),
        "majmin_wcsr": duration_weighted_accuracy(
            reference_items, estimated_items,
            lambda left, right: (
                raiz_token(left), qualidade_token(left)
            ) == (
                raiz_token(right), qualidade_token(right)
            ),
        ),
        "full_quality_wcsr": duration_weighted_accuracy(
            reference_items, estimated_items,
            lambda left, right: (
                raiz_token(left), qualidade_completa_token(left)
            ) == (
                raiz_token(right), qualidade_completa_token(right)
            ),
        ),
        "no_chord_wcsr": duration_weighted_accuracy(
            reference_items, estimated_items,
            lambda left, right: eh_sem_acorde(left) == eh_sem_acorde(right),
        ),
        "boundaries": boundary_f1(reference_items, estimated_items, tolerance_seconds),
    }


def label_at(intervals: list[ChordInterval], timestamp: float) -> str:
    for item in intervals:
        if item.start <= timestamp < item.end:
            return item.label
    return "N"


def internal_boundaries(intervals: list[ChordInterval]) -> list[float]:
    return [item.end for item in intervals[:-1]]