#!/usr/bin/env python3 """Figures for the write-up (reads data/results.json, data/est_voice.wav; re-runs pitch tracking).""" import json import os import sys import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt # noqa: E402 import numpy as np # noqa: E402 import soundfile as sf # noqa: E402 HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) import run as R # noqa: E402 import unmix as U # noqa: E402 import voiceline as VL # noqa: E402 INK, MUTED, RULE = "#0E1A2E", "#5A6577", "#D8DDE6" ACCENT, GOLD, TEAL, RED = "#2153D0", "#C98A00", "#0E8C7A", "#C2412D" plt.rcParams.update({"font.size": 11, "axes.edgecolor": RULE, "axes.labelcolor": INK, "xtick.color": MUTED, "ytick.color": MUTED, "axes.grid": True, "grid.color": "#EEF1F5", "axes.axisbelow": True, "figure.dpi": 150, "savefig.bbox": "tight", "axes.spines.top": False, "axes.spines.right": False}) OUT = os.path.join(HERE, "out") os.makedirs(OUT, exist_ok=True) def excerpt(d, start=2.5, secs=20.0): g_all = R.load_mp3(os.path.join(d, "mix.mp3")) st = {k: sf.read(os.path.join(d, "stems", k + ".wav"), dtype="float32")[0].T for k in ["voice", "chords", "bass", "drums"]} s0, T = int(start * R.SR), (int(secs * R.SR) // 256) * 256 cut = lambda x: x[..., s0:s0 + T] return cut(g_all), {k: cut(v) for k, v in st.items()}, cut(np.load(os.path.join(d, "f0_true.npy"))), T def fig_results(res): order = [("mixture", "Mixture (no separation)"), ("repet-sim", "REPET-SIM (baseline)"), ("comb-mask", "Comb mask"), ("pdhg-naive-f0", "PDHG, plain pitch tracker"), ("pdhg-mono", "PDHG, ignoring stereo"), ("pdhg", "PDHG"), ("pdhg+refit", "PDHG + refit (proposed)"), ("comb-mask-oracle", "Comb mask, true pitch"), ("pdhg-oracle-f0+refit", "PDHG + refit, true pitch")] order = [(k, n) for k, n in order if k in res][::-1] fig, ax = plt.subplots(1, 2, figsize=(11, 4.3), sharey=True) for a, (key, title) in zip(ax, [("voice_si_sdr", "Voice, SI-SDR (dB)"), ("acc_sdr", "Accompaniment, SDR (dB)")]): vals = [res[k][key] for k, _ in order] cols = [ACCENT if k == "pdhg+refit" else ("#5A6577" if "oracle" in k else "#A9BCEB") for k, _ in order] y = np.arange(len(order)) a.barh(y, vals, color=cols, height=0.66) for yi, v in zip(y, vals): a.text(v + (0.25 if v >= 0 else -0.25), yi, f"{(0.0 if abs(v) < 0.05 else v):.1f}", va="center", ha="left" if v >= 0 else "right", fontsize=9, color=INK) a.set_yticks(y, [n for _, n in order]) a.axvline(0, color=MUTED, lw=0.8) a.set_title(title, fontsize=12, color=INK, loc="left") a.grid(axis="y", visible=False) lo, hi = min(vals + [0]), max(vals) a.set_xlim(lo - 1.6 if lo < 0 else 0, hi + 2.2) fig.savefig(os.path.join(OUT, "results.png")) plt.close(fig) def fig_tracking(): rows = [] for d, name in [("data", "evaluation song (alto)"), ("data_tune", "tuning song (tenor)")]: g, st, f0s, T = excerpt(os.path.join(HERE, d)) F = U.Frame(T) f0t = R.oracle_f0(f0s, F) steps = [("plain tracker", dict(centre_power=0, vibrato=0, voicing=0.25, fold=False)), ("+ centre", dict(centre_power=80, vibrato=0, voicing=0.25, fold=False)), ("+ vibrato", dict(centre_power=80, vibrato=3, voicing=0.25, fold=False)), ("+ voicing", dict(centre_power=80, vibrato=3, voicing=0.03, fold=False)), ("+ register", dict(centre_power=80, vibrato=3, voicing=0.03, fold=True))] rows.append((name, [(s, R.rpa(VL.track(g, F, **kw), f0t)) for s, kw in steps])) print(name, [(s, round(v, 3)) for s, v in rows[-1][1]], flush=True) fig, ax = plt.subplots(figsize=(8.5, 3.6)) w = 0.38 labels = [s for s, _ in rows[0][1]] for i, (name, vals) in enumerate(rows): x = np.arange(len(labels)) + (i - 0.5) * w b = ax.bar(x, [100 * v for _, v in vals], w, color=[ACCENT, GOLD][i], label=name) for rect, (_, v) in zip(b, vals): ax.text(rect.get_x() + rect.get_width() / 2, 100 * v + 1, f"{100 * v:.0f}", ha="center", fontsize=8.5, color=INK) ax.set_xticks(range(len(labels)), labels) ax.set_ylabel("voice pitch correct (%)") ax.set_ylim(0, 80) ax.legend(frameon=False, loc="upper left") ax.grid(axis="x", visible=False) fig.savefig(os.path.join(OUT, "tracking.png")) plt.close(fig) return rows def fig_spectrograms(): g, st, f0s, T = excerpt(os.path.join(HERE, "data")) est = sf.read(os.path.join(HERE, "data", "est_voice.wav"), dtype="float32")[0] F = U.Frame(T) f0_hat = VL.track(g, F) f0t = R.oracle_f0(f0s, F) t = F.frame_times() a, b = np.searchsorted(t, 4.0), np.searchsorted(t, 10.0) hz = F.bin_hz() kmax = int(3000 / (R.SR / F.N)) panels = [("Mixture (centre of the stereo mix)", g.mean(0)), ("True voice", st["voice"][0]), ("Recovered voice (PDHG + refit)", est)] fig, ax = plt.subplots(len(panels), 1, figsize=(10, 7.4), sharex=True) ref = None for i, (title, x) in enumerate(panels): S = 20 * np.log10(np.abs(F.fwd(x))[a:b, :kmax].T + 1e-6) ref = ref if ref is not None else S.max() ax[i].imshow(S, origin="lower", aspect="auto", cmap="magma", vmin=ref - 70, vmax=ref, extent=[t[a], t[b - 1], 0, hz[kmax - 1]]) ax[i].set_title(title, fontsize=11, loc="left", color=INK) ax[i].set_ylabel("Hz") ax[i].grid(False) if i == 0: tt = t[a:b] ax[i].plot(tt, np.where(f0t[a:b] > 0, f0t[a:b], np.nan), color="#FFFFFF", lw=1.2, label="true pitch") ax[i].plot(tt, np.where(f0_hat[a:b] > 0, f0_hat[a:b], np.nan), color="#4FD1C5", lw=1.2, ls="--", label="tracked pitch") ax[i].legend(loc="upper right", fontsize=8, frameon=False, labelcolor="#FFFFFF") ax[-1].set_xlabel("seconds") fig.savefig(os.path.join(OUT, "spectrograms.png")) plt.close(fig) if __name__ == "__main__": res = json.load(open(os.path.join(HERE, "data", "results.json")) if os.path.exists(os.path.join(HERE, "data", "results.json")) else open(os.path.join(HERE, "results", "eval_song.json"))) fig_results(res["results"]) fig_spectrograms() rows = fig_tracking() json.dump({n: dict(v) for n, v in rows}, open(os.path.join(OUT, "tracking.json"), "w"), indent=1) print("figures in", OUT)