ensemble: merge_transcriptions.py

merge_transcriptions.py

@@ -0,0 +1,251 @@
+#!/usr/bin/env python3
+"""Merge several MuScriptor passes over the same audio into one transcription.
+
+A single pass of a transcription model makes two kinds of mistake, and they are
+not the same mistake twice: one pass drops the blurry note in a fast run, another
+pass invents a ghost in a reverb tail. Running the model a few times — different
+sizes, different instrument conditioning — and keeping only what the passes agree
+on removes both, because a real note is found again and a hallucination is not.
+
+This is the one "global" improvement available without a GPU on the training
+side: no weights change, no extra data, only a second opinion.
+
+Voting rules:
+  * notes cluster by pitch and onset within `tolerance`;
+  * a cluster survives when at least `votes` distinct passes found it;
+  * its onset is the **median** of the passes, which is steadier than any one;
+  * velocity is measured from the audio at the agreed onset, since no pass
+    predicts loudness.
+
+Usage:
+    python3 merge_transcriptions.py song.wav out.mid pass1.mid pass2.mid pass3.mid
+    python3 merge_transcriptions.py song.wav out.mid --votes 2 --keep-strong a.mid b.mid
+    python3 merge_transcriptions.py --selftest
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import sys
+
+import numpy as np
+
+from muscriptor_dynamics import load_audio, measure, read_midi, write_midi
+
+
+def _median(values):
+    ordered = sorted(values)
+    middle = len(ordered) // 2
+    if len(ordered) % 2:
+        return ordered[middle]
+    return 0.5 * (ordered[middle - 1] + ordered[middle])
+
+
+def cluster(passes, tolerance=0.060):
+    """Group notes that the passes agree on.
+
+    `passes` is a list of note lists, each note (start, dur, pitch, channel).
+    Returns clusters as dicts with the members' pass indices.
+    """
+    flat = []
+    for index, notes in enumerate(passes):
+        for note in notes:
+            flat.append((note[0], note[1], note[2], note[3], index))
+    flat.sort()
+
+    clusters = []
+    for start, dur, pitch, channel, which in flat:
+        placed = False
+        for item in reversed(clusters[-8:]):
+            if item["pitch"] != pitch or which in item["passes"]:
+                continue
+            if abs(item["times"][0] - start) > tolerance:
+                continue
+            item["times"].append(start)
+            item["durs"].append(dur)
+            item["passes"].add(which)
+            item["channel"] = channel
+            placed = True
+            break
+        if not placed:
+            clusters.append({"pitch": pitch, "channel": channel, "times": [start],
+                             "durs": [dur], "passes": {which}})
+    return clusters
+
+
+def merge(passes, samples=None, rate=None, votes=2, tolerance=0.060,
+          keep_strong_singletons=False, on_stage=None):
+    """One transcription out of several passes. Returns (notes, report)."""
+    def stage(percent, message):
+        if on_stage:
+            on_stage(percent, message)
+
+    stage(5, "группирую ноты из %d проходов" % len(passes))
+    clusters = cluster(passes, tolerance)
+    stage(35, "согласовано кластеров: %d" % len(clusters))
+
+    kept = [item for item in clusters if len(item["passes"]) >= votes]
+    report = {
+        "passes": len(passes),
+        "clusters": len(clusters),
+        "agreed": len(kept),
+        "dropped_doubtful": len(clusters) - len(kept),
+        "rescued": 0,
+    }
+    stage(55, "прошли голосование: %d" % len(kept))
+
+    if keep_strong_singletons and samples is not None and rate:
+        # A note that only one pass heard is usually a hallucination, but a loud,
+        # clean attack is not: a fortissimo note in an otherwise quiet passage is
+        # heard once and is real. Those are rescued by energy, not by vote.
+        levels = []
+        for item in kept:
+            first = int(max(0, round(_median(item["times"]) * rate)))
+            chunk = samples[first:first + max(1, int(0.06 * rate))]
+            levels.append(20 * math.log10(float(np.sqrt(np.mean(chunk * chunk)) + 1e-9))
+                          if chunk.size else -120.0)
+        floor = float(np.percentile(levels, 10)) if levels else -60.0
+        for item in clusters:
+            if len(item["passes"]) >= votes:
+                continue
+            first = int(max(0, round(_median(item["times"]) * rate)))
+            chunk = samples[first:first + max(1, int(0.06 * rate))]
+            if not chunk.size:
+                continue
+            level = 20 * math.log10(float(np.sqrt(np.mean(chunk * chunk)) + 1e-9))
+            if level > floor + 6.0:
+                kept.append(item)
+                report["rescued"] += 1
+        kept.sort(key=lambda item: _median(item["times"]))
+        stage(65, "спасено по атаке: %d" % report["rescued"])
+
+    notes = [(_median(item["times"]), max(0.05, _median(item["durs"])),
+              item["pitch"], item["channel"]) for item in kept]
+    notes.sort(key=lambda note: note[0])
+    stage(75, "слито нот: %d" % len(notes))
+
+    if samples is not None and rate and notes:
+        stage(85, "меряю velocity по аудио")
+        velocities = measure(notes, samples, rate)
+        notes = [(start, dur, pitch, channel, velocities[i])
+                 for i, (start, dur, pitch, channel) in enumerate(notes)]
+    else:
+        notes = [(start, dur, pitch, channel, 80) for start, dur, pitch, channel in notes]
+
+    report["notes"] = len(notes)
+    stage(95, "готово")
+    return notes, report
+
+
+# --------------------------------------------------------------------------- #
+# scoring, shared with transcript_fix's benchmark idea
+# --------------------------------------------------------------------------- #
+def score(reference, candidate, tolerance=0.05):
+    used = [False] * len(candidate)
+    matched = 0
+    errors = []
+    for start, _dur, pitch, *_rest in reference:
+        best = None
+        for index, other in enumerate(candidate):
+            if used[index] or other[2] != pitch:
+                continue
+            delta = abs(other[0] - start)
+            if delta <= tolerance and (best is None or delta < best[0]):
+                best = (delta, index)
+        if best is not None:
+            used[best[1]] = True
+            matched += 1
+            errors.append(best[0])
+    precision = matched / max(1, len(candidate))
+    recall = matched / max(1, len(reference))
+    return {
+        "precision": precision,
+        "recall": recall,
+        "f1": 2 * precision * recall / max(1e-9, precision + recall),
+        "mean_error_ms": 1000.0 * float(np.mean(errors)) if errors else float("nan"),
+    }
+
+
+def selftest():
+    """Three simulated passes, each losing different notes.
+
+    This is the case ensembles exist for: the union of what they miss must not
+    survive, and the union of what they find must.
+    """
+    truth = [
+        (0.00, 0.24, 60), (0.30, 0.22, 62), (0.58, 0.20, 64), (0.86, 0.26, 65),
+        (1.14, 0.22, 67), (1.42, 0.24, 69), (1.70, 0.20, 71), (1.98, 0.26, 72),
+        (2.26, 0.22, 74), (2.54, 0.20, 75), (2.82, 0.24, 76), (3.10, 0.26, 77),
+        (3.38, 0.22, 79), (3.66, 0.24, 81),
+    ]
+    reference = [(start, dur, pitch, 0, 90) for start, dur, pitch in truth]
+
+    missing = ({2, 7, 11}, {4, 9, 12}, {1, 6, 10})
+    jitter = (0.010, 0.012, 0.008)
+    ghosts = (3.98, 2.40, 0.44)
+    rng = np.random.default_rng(11)
+
+    passes = []
+    for index in range(3):
+        notes = []
+        for position, (start, dur, pitch) in enumerate(truth):
+            if position in missing[index]:
+                continue
+            notes.append((max(0.0, start + rng.normal(0, jitter[index])), dur, pitch, 0))
+        notes.append((ghosts[index], 0.2, 55 + index, 0))
+        notes.sort()
+        passes.append(notes)
+        result = score(reference, [(n[0], n[1], n[2], n[3], 80) for n in notes])
+        print("проход %d: P=%.2f R=%.2f F1=%.2f, нот %d"
+              % (index + 1, result["precision"], result["recall"], result["f1"], len(notes)),
+              flush=True)
+
+    merged, report = merge(passes, votes=2,
+                           on_stage=lambda p, m: print("  [%3d%%] %s" % (p, m), flush=True))
+    result = score(reference, merged)
+    print("слияние:  P=%.2f R=%.2f F1=%.2f, ошибка %.1f мс, нот %d"
+          % (result["precision"], result["recall"], result["f1"],
+             result["mean_error_ms"], len(merged)), flush=True)
+    print("отчёт: %s" % report, flush=True)
+
+    best_single = max(score(reference, [(n[0], n[1], n[2], n[3], 80) for n in notes])["f1"]
+                      for notes in passes)
+    assert result["f1"] > best_single + 0.15, "ансамбль не лучше лучшего прохода"
+    assert result["recall"] > 0.95, "потерянные ноты не восстановлены"
+    assert result["precision"] > 0.95, "фантомы не отфильтрованы"
+    print("selftest: OK", flush=True)
+
+
+def run(audio_path, out_path, midi_paths, votes=2, keep_strong=False):
+    print("[  0%] читаю аудио", flush=True)
+    samples, rate = load_audio(audio_path)
+    passes = []
+    for path in midi_paths:
+        notes, _tempo, _division = read_midi(path)
+        passes.append([note[:4] for note in notes])
+        print("        %s: %d нот" % (os.path.basename(path), len(notes)), flush=True)
+    if len(passes) < votes:
+        raise SystemExit("проходов меньше, чем нужно голосов")
+    merged, report = merge(passes, samples, rate, votes=votes,
+                           keep_strong_singletons=keep_strong,
+                           on_stage=lambda p, m: print("[%3d%%] %s" % (p, m), flush=True))
+    write_midi(out_path, merged, 120.0)
+    print("[100%] %s: %s" % (out_path, report), flush=True)
+    return merged
+
+
+if __name__ == "__main__":
+    if "--selftest" in sys.argv:
+        selftest()
+    elif len(sys.argv) >= 4:
+        args = [a for a in sys.argv[1:] if not a.startswith("--")]
+        keep = "--keep-strong" in sys.argv
+        votes = 2
+        if "--votes" in sys.argv:
+            votes = int(sys.argv[sys.argv.index("--votes") + 1])
+            args = [a for a in args if a != str(votes)]
+        run(args[0], args[1], args[2:], votes=votes, keep_strong=keep)
+    else:
+        print(__doc__)
+        raise SystemExit(2)