Audio compressors on MUSDB

Side-by-side rate-distortion plots, sourced from the JSONs hardcoded below. All metrics come from rate_distortion_<id>.json measured on the 262 stereo 44.1 kHz clips of the danjacobellis/musdb_segments validation split: PSNR on [0, 1]-convention waveforms, and SSDR / SRDR via compressors.spatial_audio_quality (Watcharasupat & Lerch, ICASSP 2024). Rate is reported in bits per sample of the 44.1 kHz timeline (channels combined), mirroring the bits-per-pixel convention on the image side.

WaLLoC, LiVeAction, and the Stable Audio VAE are single-operating-point codecs and appear as lone markers; the others sweep a rate knob. No encode-throughput panel yet — the audio complexity harnesses exist but a rigorous quiet-machine sweep has not been run.

Pick which codecs to plot by commenting / uncommenting lines in the CODECS list below.

import json
from pathlib import Path
import matplotlib.pyplot as plt
plt.rcParams["font.family"] = "serif"

RD_PATHS = {
    "mp3":              "results/mp3/rate_distortion_1782553179.json",
    "opus":             "results/opus/rate_distortion_1782554683.json",
    "encodec":          "results/encodec/rate_distortion_1782556019.json",
    "dac":              "results/dac/rate_distortion_1782556377.json",
    "mimi":             "results/mimi/rate_distortion_1782556904.json",
    "oobleck":          "results/oobleck/rate_distortion_1782557202.json",
    "walloc_audio":     "results/walloc_audio/rate_distortion_1782557303.json",
    "liveaction_audio": "results/liveaction_audio/rate_distortion_1782557364.json",
    "frappe_v2_audio":  "results/frappe_v2_audio/rate_distortion_1782557426.json",
}

AXIS_LABEL = {
    "bits_per_sample": "Rate [bits per sample]",
    "PSNR_dB":         "PSNR [dB]",
    "SSDR":            "SSDR [dB]",
    "SRDR":            "SRDR [dB]",
}
TITLE_LABEL = {
    "bits_per_sample": "Rate",
    "PSNR_dB":         "PSNR",
    "SSDR":            "SSDR",
    "SRDR":            "SRDR",
}
LOG_KEYS = ("bits_per_sample",)

def load_codec(name):
    rd = json.loads(Path(RD_PATHS[name]).read_text())
    fs = rd["sample_rate"]
    points = []
    for k in rd["quality_values"]:
        rm = rd["results"][str(k)]["mean"]
        points.append({
            "sweep": k,
            "bits_per_sample": rm["kbps"] * 1000.0 / fs,
            "PSNR_dB": rm["PSNR_dB"],
            "SSDR": rm["SSDR"],
            "SRDR": rm["SRDR"],
        })
    return points

def col(pts, key):
    return [p[key] for p in pts]

def plot_panel(group, y_key, x_key, legend=True):
    """One small panel: y vs x for every codec in `group`.

    The rate axis is log-scaled wherever it appears.
    """
    plt.figure(figsize=(4.5, 4), dpi=180)
    ax = plt.gca()
    if x_key in LOG_KEYS:
        ax.set_xscale("log")
    if y_key in LOG_KEYS:
        ax.set_yscale("log")
    for c in group:
        ax.plot(col(c["data"], x_key), col(c["data"], y_key),
                marker=c.get("marker", "."),
                linestyle=c.get("linestyle", "-"),
                color=c.get("color"), label=c["name"])
    ax.set_xlabel(AXIS_LABEL[x_key])
    ax.set_ylabel(AXIS_LABEL[y_key])
    ax.set_title(f"{TITLE_LABEL[y_key]} vs {TITLE_LABEL[x_key]} (MUSDB)")
    if legend:
        ax.legend(loc="best")
    log_either = (x_key in LOG_KEYS) or (y_key in LOG_KEYS)
    ax.grid(True, which="both" if log_either else "major", alpha=0.4)
    plt.tight_layout()
    plt.show()

# Comment / uncomment to control which codecs appear on the plots
CODECS = [
    {"name": "MP3",              "data": load_codec("mp3"),              "color": "gray",       "linestyle": "-", "marker": "."},
    {"name": "Opus",             "data": load_codec("opus"),             "color": "black",      "linestyle": "-", "marker": "."},
    {"name": "EnCodec",          "data": load_codec("encodec"),          "color": "red",        "linestyle": "-", "marker": "."},
    {"name": "DAC",              "data": load_codec("dac"),              "color": "firebrick",  "linestyle": "-", "marker": "."},
    {"name": "Mimi",             "data": load_codec("mimi"),             "color": "orange",     "linestyle": "-", "marker": "."},
    {"name": "Stable Audio VAE", "data": load_codec("oobleck"),          "color": "tab:cyan",   "linestyle": "none", "marker": "D"},
    {"name": "WaLLoC",           "data": load_codec("walloc_audio"),     "color": "blue",       "linestyle": "none", "marker": "o"},
    {"name": "LiVeAction",       "data": load_codec("liveaction_audio"), "color": "green",      "linestyle": "none", "marker": "s"},
    {"name": "FRAPPE v2",        "data": load_codec("frappe_v2_audio"),  "color": "magenta",    "linestyle": "-",    "marker": "."},
]

for c in CODECS:
    bps = col(c["data"], "bits_per_sample")
    print(f"{c['name']:>20}: {len(c['data'])} points,"
          f" bits/sample [{min(bps):.4f}, {max(bps):.4f}],"
          f" PSNR [{min(col(c['data'],'PSNR_dB')):.2f}, {max(col(c['data'],'PSNR_dB')):.2f}] dB")

                 MP3: 10 points, bits/sample [0.7272, 5.8133], PSNR [34.73, 53.31] dB
                Opus: 10 points, bits/sample [0.1762, 5.8483], PSNR [28.07, 47.60] dB
             EnCodec: 5 points, bits/sample [0.0680, 1.0886], PSNR [26.19, 32.20] dB
                 DAC: 5 points, bits/sample [0.0227, 0.2720], PSNR [23.31, 24.22] dB
                Mimi: 6 points, bits/sample [0.0062, 0.1997], PSNR [21.72, 29.44] dB
    Stable Audio VAE: 1 points, bits/sample [0.5000, 0.5000], PSNR [28.41, 28.41] dB
              WaLLoC: 1 points, bits/sample [1.4969, 1.4969], PSNR [39.02, 39.02] dB
          LiVeAction: 1 points, bits/sample [0.1090, 0.1090], PSNR [31.07, 31.07] dB
           FRAPPE v2: 7 points, bits/sample [0.0360, 3.6565], PSNR [27.96, 42.49] dB

plot_panel(CODECS, "PSNR_dB", "bits_per_sample")

plot_panel(CODECS, "SSDR", "bits_per_sample", legend=False)

plot_panel(CODECS, "SRDR", "bits_per_sample", legend=False)