docs/make_figures.py10.5 KB
"""Generate the figures used in the README.
Run from the project root with the package importable::
python docs/make_figures.py
Writes PNGs into ``docs/figures/``. The plots are produced from the real pipeline and the
synthetic fixtures, so they stay in sync with the implementation.
"""
from __future__ import annotations
import sys
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "tests"))
from localizer.audio_io import StereoIR # noqa: E402
from localizer.config import ExtractionConfig, GammatoneConfig # noqa: E402
from localizer.extraction import extract_hrir, _analytic_envelope, _smooth # noqa: E402
from localizer.gammatone import GammatoneFilterbank # noqa: E402
from localizer.cues import compute_cues # noqa: E402
from fixtures import synth_brir, synth_hrir, SAMPLE_RATE # noqa: E402
FIG = ROOT / "docs" / "figures"
FIG.mkdir(parents=True, exist_ok=True)
# A calm, readable style that renders well framed on a dark site or on GitHub.
plt.rcParams.update({
"figure.facecolor": "white",
"axes.facecolor": "white",
"font.size": 11,
"axes.titlesize": 12,
"axes.grid": True,
"grid.alpha": 0.25,
"savefig.bbox": "tight",
"savefig.dpi": 130,
})
INK = "#1b2b34"
BLUE = "#2a6f97"
RED = "#c1121f"
AMBER = "#e09f3e"
GREEN = "#2a9d8f"
def fig_pipeline():
fig, ax = plt.subplots(figsize=(11, 2.6))
ax.axis("off")
ax.set_xlim(0, 100)
ax.set_ylim(0, 24)
stages = [
("BRIR", "two-channel\nroom response", "#e8eef2"),
("HRIR\nextraction", "6 ms direct\nsound window", "#d6e4ee"),
("Cue\nextraction", "ITD · ILD\ngammatone spectrum", "#c3d9e8"),
("Compare vs\nreference DB", "per subject &\ndirection", "#b0cfe2"),
("Azimuth", "lateral angle\n+ front/back", "#e09f3e"),
]
x = 1.5
w, h, gap = 15, 12, 4.25
centers = []
for title, sub, color in stages:
box = FancyBboxPatch((x, 7), w, h, boxstyle="round,pad=0.3,rounding_size=0.8",
linewidth=1.2, edgecolor=INK, facecolor=color)
ax.add_patch(box)
ax.text(x + w / 2, 7 + h / 2 + 1.6, title, ha="center", va="center",
fontweight="bold", color=INK)
ax.text(x + w / 2, 7 + h / 2 - 2.6, sub, ha="center", va="center",
fontsize=8.5, color="#425a66")
centers.append(x + w)
x += w + gap
for cx in centers[:-1]:
ax.add_patch(FancyArrowPatch((cx + 0.3, 13), (cx + gap - 0.3, 13),
arrowstyle="-|>", mutation_scale=16,
linewidth=1.6, color=INK))
ax.text(50, 2, "Model of human binaural hearing, applied directly to HRTFs",
ha="center", fontsize=9, style="italic", color="#607d8b")
fig.savefig(FIG / "pipeline.png")
plt.close(fig)
def fig_extraction():
brir = synth_brir(35.0, leading_silence=60, tail_len=2200)
cfg = ExtractionConfig()
hrir = extract_hrir(brir, cfg)
sr = brir.sample_rate
summed = brir.left + brir.right
env = _smooth(_analytic_envelope(summed), max(1, int(0.0005 * sr)))
n_show = 900
t = np.arange(n_show) / sr * 1000.0
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 3.6),
gridspec_kw={"width_ratios": [2, 1]})
ax1.plot(t, brir.left[:n_show], color=BLUE, lw=0.7, label="left")
ax1.plot(t, brir.right[:n_show], color=AMBER, lw=0.7, alpha=0.8, label="right")
ax1.plot(t, env[:n_show] / env.max() * np.abs(brir.left[:n_show]).max(),
color=RED, lw=1.4, label="envelope")
onset = int(np.argmax(env))
win = int(cfg.hrir_duration * sr)
ax1.axvspan(onset / sr * 1000, (onset + win) / sr * 1000, color=GREEN, alpha=0.15)
ax1.axvline(onset / sr * 1000, color=GREEN, lw=1.2, ls="--")
ax1.text(onset / sr * 1000 + 0.15, ax1.get_ylim()[1] * 0.8,
"6 ms window", color=GREEN, fontsize=9)
ax1.set_xlabel("time (ms)")
ax1.set_ylabel("amplitude")
ax1.set_title("BRIR: locate direct sound, take a 6 ms window")
ax1.legend(loc="upper right", fontsize=8)
th = np.arange(hrir.num_samples) / sr * 1000.0
ax2.plot(th, hrir.left, color=BLUE, lw=0.9, label="left")
ax2.plot(th, hrir.right, color=AMBER, lw=0.9, label="right")
ax2.set_xlabel("time (ms)")
ax2.set_title("Extracted HRIR")
ax2.legend(loc="upper right", fontsize=8)
fig.savefig(FIG / "hrir_extraction.png")
plt.close(fig)
def fig_gammatone():
sr = 44100
fb = GammatoneFilterbank(GammatoneConfig(), sr)
n = 8192
imp = np.zeros(n)
imp[0] = 1.0
freqs = np.fft.rfftfreq(n, 1 / sr)
fig, ax = plt.subplots(figsize=(11, 3.8))
for i, band in enumerate(fb.filter(imp)):
H = np.abs(np.fft.rfft(band))
H_db = 20 * np.log10(H + 1e-9)
color = plt.cm.viridis(i / (len(fb.center_freqs) - 1))
ax.plot(freqs, H_db, color=color, lw=1.0)
ax.set_xscale("log")
ax.set_xlim(400, 22000)
ax.set_ylim(-45, 3)
ax.set_xlabel("frequency (Hz)")
ax.set_ylabel("magnitude (dB)")
ax.set_title("28-band ERB-spaced gammatone filterbank (700 Hz – 18 kHz)")
for cf in fb.center_freqs:
ax.axvline(cf, color="k", alpha=0.05)
fig.savefig(FIG / "gammatone_filterbank.png")
plt.close(fig)
def fig_cues():
hrir = synth_hrir(50.0) # source on the left
cues = compute_cues(hrir, gammatone_config=GammatoneConfig())
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 3.6))
sr = hrir.sample_rate
t = np.arange(hrir.num_samples) / sr * 1000.0
ax1.plot(t, hrir.left, color=BLUE, lw=0.9, label="left (near ear)")
ax1.plot(t, hrir.right, color=AMBER, lw=0.9, label="right (far ear)")
ax1.set_xlim(0, 2.2)
ax1.set_xlabel("time (ms)")
ax1.set_ylabel("amplitude")
ax1.set_title(f"ITD = {cues.itd * 1000:+.2f} ms (right ear delayed)")
ax1.legend(loc="upper right", fontsize=8)
ax2.plot(cues.center_freqs, cues.ild_bands, "-o", ms=3, color=GREEN)
ax2.axhline(0, color="k", lw=0.8, alpha=0.4)
ax2.set_xscale("log")
ax2.set_xlabel("band centre frequency (Hz)")
ax2.set_ylabel("ILD, left − right (dB)")
ax2.set_title(f"Per-band ILD (mean {cues.ild_broadband:+.1f} dB)")
fig.savefig(FIG / "cues.png")
plt.close(fig)
def fig_itd_ild_vs_azimuth():
az = np.arange(-180, 181, 10)
itds, ilds = [], []
fb = GammatoneFilterbank(GammatoneConfig(), SAMPLE_RATE)
for a in az:
c = compute_cues(synth_hrir(float(a)), filterbank=fb)
itds.append(c.itd * 1000)
ilds.append(c.ild_broadband)
fig, ax = plt.subplots(figsize=(11, 3.8))
ax.plot(az, itds, "-o", ms=3, color=BLUE, label="ITD (ms)")
ax.plot(az, ilds, "-s", ms=3, color=RED, label="ILD (dB, ×0.1 scale)")
ax.set_xlabel("true azimuth (degrees) — 0° front, +90° left, ±180° back")
ax.set_ylabel("cue value")
ax.axvspan(90, 180, color="gray", alpha=0.08)
ax.axvspan(-180, -90, color="gray", alpha=0.08)
ax.set_title("Front/back ambiguity: a front azimuth and its rear mirror share ITD & ILD")
# annotate mirror pair
ax.annotate("30° and 150°\nsame ITD/ILD",
xy=(30, np.interp(30, az, itds)), xytext=(60, 0.55),
fontsize=8.5, color=INK,
arrowprops=dict(arrowstyle="->", color=INK, lw=0.8))
ax.plot([150], [np.interp(150, az, itds)], "o", color=BLUE, ms=7, mfc="none")
ax.legend(loc="upper left", fontsize=9)
# scale ILD for display
ax.lines[1].set_ydata(np.array(ilds) * 0.1)
ax.relim(); ax.autoscale_view()
fig.savefig(FIG / "itd_ild_vs_azimuth.png")
plt.close(fig)
def fig_polar():
import tempfile
from localizer.pipeline import Localizer
from localizer.config import PipelineConfig
from localizer import plotting
sys.path.insert(0, str(ROOT / "tests"))
from fixtures import synth_database_dir
tmp = Path(tempfile.mkdtemp())
db = synth_database_dir(tmp)
loc = Localizer.from_database_path(db, config=PipelineConfig())
result = loc.analyze(synth_brir(120.0))
fig = plt.figure(figsize=(11, 4.8))
ax_s = fig.add_subplot(121, projection="polar")
plotting.plot_polar_similarity(result, "specialized", ax=ax_s)
ax_g = fig.add_subplot(122, projection="polar")
plotting.plot_polar_similarity(result, "global", ax=ax_g)
fig.text(0.5, 0.02,
"True source at 120° (back-left). Specialized nails it; global misfires to 90°.",
ha="center", fontsize=9.5, style="italic", color="#607d8b")
fig.subplots_adjust(top=0.88, bottom=0.14)
fig.savefig(FIG / "polar_similarity.png")
plt.close(fig)
def fig_cnn():
fig, ax = plt.subplots(figsize=(11, 3.0))
ax.axis("off")
ax.set_xlim(0, 100)
ax.set_ylim(0, 20)
layers = [
("input", "2 × 256", "#e8eef2"),
("conv + ReLU", "16 × 256", "#d6e4ee"),
("maxpool", "16 × 128", "#c9dbe8"),
("conv + ReLU", "32 × 128", "#bcd2e3"),
("conv + ReLU", "64 × 128", "#aecadd"),
("dense + ReLU\ndropout 0.5", "128", "#a0c1d8"),
("output\nsoftmax", "front / back", "#e09f3e"),
]
x, w, h, gap = 2, 11, 12, 2.5
centers = []
for name, shape, color in layers:
box = FancyBboxPatch((x, 4), w, h, boxstyle="round,pad=0.2,rounding_size=0.6",
linewidth=1.1, edgecolor=INK, facecolor=color)
ax.add_patch(box)
ax.text(x + w / 2, 4 + h / 2 + 1.4, name, ha="center", va="center",
fontsize=8.6, fontweight="bold", color=INK)
ax.text(x + w / 2, 4 + h / 2 - 2.4, shape, ha="center", va="center",
fontsize=8, color="#425a66")
centers.append(x + w)
x += w + gap
for cx in centers[:-1]:
ax.add_patch(FancyArrowPatch((cx + 0.2, 10), (cx + gap - 0.2, 10),
arrowstyle="-|>", mutation_scale=12,
linewidth=1.3, color=INK))
ax.set_title("Front/back CNN — raw HRIR in, hemisphere probabilities out", fontsize=12)
fig.savefig(FIG / "cnn_architecture.png")
plt.close(fig)
def main():
fig_pipeline()
fig_extraction()
fig_gammatone()
fig_cues()
fig_itd_ild_vs_azimuth()
fig_polar()
fig_cnn()
print(f"Wrote figures to {FIG}")
for p in sorted(FIG.glob("*.png")):
print(" ·", p.name)
if __name__ == "__main__":
main()