Blame for plus/plus_ensemble.py

"""PLUS ensemble: слияние нескольких транскрипций голосованием.

Один проход ошибается по-разному: MuScriptor теряет размытые ноты и добавляет
фантомы, Aria путает октавы, третий движок молчит на группе. Слияние берёт то, в чём
сошлись РАЗНЫЕ движки: кластер по питчу и онсету, выживает при сумме весов
источников >= порога.

CLI:
  python plus_ensemble.py --inputs m.mid aria.mid other.mid --out merged.mid \
      --tolerance 0.06 --min-votes 2 --reference truth.mid
"""
from __future__ import annotations

import argparse
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from plus_core import Note, Piece, dump_json, load_midi, save_midi, step  # noqa: E402

__version__ = "1.0"


def _cluster(notes: list[Note], tolerance: float) -> list[list[int]]:
    """Индексы нот, сгруппированные по (питч, онсет в пределах допуска)."""
    order = sorted(range(len(notes)), key=lambda i: (notes[i].pitch, notes[i].start))
    clusters: list[list[int]] = []
    for index in order:
        note = notes[index]
        placed = False
        for cluster in reversed(clusters[-6:]):
            head = notes[cluster[0]]
            if head.pitch != note.pitch:
                continue
            if abs(head.start - note.start) <= tolerance:
                cluster.append(index)
                placed = True
                break
        if not placed:
            clusters.append([index])
    return clusters


def merge(sources: list[tuple[str, list[Note], float]], tolerance: float,
          min_votes: float, keep_strong: bool, audio: str | None = None,
          tempo: float | None = None) -> tuple[Piece, dict]:
    """Сливает источники и возвращает Piece плюс статистику согласия."""
    flattened: list[Note] = []
    owner: list[str] = []
    for name, notes, _weight in sources:
        for note in notes:
            flattened.append(note)
            owner.append(name)
    if not flattened:
        return Piece(), {"notes": 0}

    weight_of = {name: weight for name, _notes, weight in sources}
    clusters = _cluster(flattened, tolerance)
    merged: list[Note] = []
    rejected = 0
    agreements: list[int] = []
    for cluster in clusters:
        per_source: dict[str, list[Note]] = {}
        for index in cluster:
            per_source.setdefault(owner[index], []).append(flattened[index])
        votes = sum(weight_of[name] for name in per_source)
        airs = [n for notes_ in per_source.values() for n in notes_]
        strong = max((n.velocity for n in airs), default=0) >= 100 and max(
            (n.dur for n in airs), default=0) >= 0.08
        if votes < min_votes and not (keep_strong and strong):
            rejected += 1
            continue
        agreements.append(len(per_source))
        starts = sorted(n.start for n in airs)
        durs = sorted(n.dur for n in airs)
        velocities = sorted(n.velocity for n in airs)
        merged.append(Note(
            start=float(np.median(starts)),
            dur=float(np.median(durs)),
            pitch=int(sorted(n.pitch for n in airs)[len(airs) // 2]),
            velocity=int(np.median(velocities)),
            channel=0,
            track=airs[0].track,
        ))
    merged.sort(key=lambda n: (n.start, n.pitch))
    piece = Piece(merged, [(0.0, tempo if tempo else 120.0)])
    piece.meta["sources"] = [name for name, _n, _w in sources]
    stats = {
        "clusters": len(clusters), "kept": len(merged), "rejected": rejected,
        "agreement_mean": round(float(np.mean(agreements)), 2) if agreements else 0.0,
        "agreement_histogram": {str(k): agreements.count(k) for k in sorted(set(agreements))},
        "per_source": {name: len(notes) for name, notes, _w in sources},
    }
    return piece, stats


def compare(reference: list[Note], candidate: list[Note], tolerance: float = 0.05) -> dict:
    """Точность/полнота/F1 по совпадению питча и онсета."""
    used = set()
    hits = 0
    errors = []
    for ref in sorted(reference, key=lambda n: n.start):
        best, best_gap = None, 1e9
        for index, cand in enumerate(candidate):
            if index in used or cand.pitch != ref.pitch:
                continue
            gap = abs(cand.start - ref.start)
            if gap < best_gap:
                best, best_gap = index, gap
        if best is not None and best_gap <= tolerance:
            used.add(best)
            hits += 1
            errors.append(best_gap)
    precision = hits / max(1, len(candidate))
    recall = hits / max(1, len(reference))
    f1 = 2 * precision * recall / max(1e-9, precision + recall)
    # numpy-скаляры не сериализуются в JSON, поэтому всё приводим к питоновским
    return {
        "precision": float(round(float(precision), 3)),
        "recall": float(round(float(recall), 3)),
        "f1": float(round(float(f1), 3)),
        "onset_error_ms": float(round(1000 * float(np.mean(errors)), 1)) if errors else None,
        "reference": int(len(reference)), "candidate": int(len(candidate)),
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="PLUS ensemble: слияние транскрипций")
    parser.add_argument("--inputs", nargs="+", required=True, help="MIDI от разных движков/проходов")
    parser.add_argument("--weights", nargs="*", type=float, help="вес на каждый вход")
    parser.add_argument("--out", required=True)
    parser.add_argument("--tolerance", type=float, default=0.06)
    parser.add_argument("--min-votes", type=float, default=2.0)
    parser.add_argument("--keep-strong", action="store_true", help="спасать одиночные громкие ноты")
    parser.add_argument("--reference", help="эталонный MIDI для оценки качества")
    parser.add_argument("--report")
    args = parser.parse_args(argv)

    weights = args.weights or [1.0] * len(args.inputs)
    if len(weights) != len(args.inputs):
        raise SystemExit("весов должно быть столько же, сколько входов")

    sources = []
    for index, path in enumerate(args.inputs, 1):
        piece = load_midi(path)
        name = os.path.splitext(os.path.basename(path))[0]
        step(index, len(args.inputs) + 2, "%s: %d нот" % (name, len(piece.notes)))
        sources.append((name, piece.notes, float(weights[index - 1])))

    piece, stats = merge(sources, args.tolerance, args.min_votes, args.keep_strong)
    step(len(args.inputs) + 1, len(args.inputs) + 2,
         "слияние: %d нот из %d кластеров" % (len(piece.notes), stats.get("clusters", 0)))
    save_midi(piece, args.out)

    report = {"version": __version__, "stats": stats, "tolerance": args.tolerance,
              "min_votes": args.min_votes, "keep_strong": args.keep_strong}
    if args.reference:
        reference = load_midi(args.reference).notes
        report["versus_reference"] = compare(reference, piece.notes)
        for name, notes, _weight in sources:
            report.setdefault("sources_versus_reference", {})[name] = compare(reference, notes)
    dump_json(report, args.report or os.path.splitext(args.out)[0] + "_ensemble.json")
    step(len(args.inputs) + 2, len(args.inputs) + 2, "готово: " + args.out)

    if args.reference:
        merged_f1 = report["versus_reference"]["f1"]
        best_source = max((entry["f1"], name) for name, entry
                          in report["sources_versus_reference"].items())
        print("      F1 слияния %.3f против лучшего входа %.3f (%s)"
              % (merged_f1, best_source[0], best_source[1]), flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())