diff --git a/apps/backend/ARCHITECTURE.md b/apps/backend/ARCHITECTURE.md index c624b0c9..8c1c06c4 100644 --- a/apps/backend/ARCHITECTURE.md +++ b/apps/backend/ARCHITECTURE.md @@ -74,7 +74,7 @@ Custom routes: - `GET /api/analysis-runs/{run_id}/artifacts` and `…/artifacts/{artifact_id}` - `GET /api/analysis-runs/{run_id}/source-audio` — re-serves the original ingested audio for the run. Owner-only (no admin bypass). Saves a round-trip vs looking up the source artifact id first. - `GET /api/analysis-runs/{run_id}/export/csv/{field_path}` — CSV export of a Phase 1 time-series field. See [`docs/adr/0001-phase1-json-schema-v1.md`](../../docs/adr/0001-phase1-json-schema-v1.md) and the registry in [`csv_export.py`](csv_export.py). -- `POST /api/analysis-runs/{run_id}/spectral-enhancements/{kind}` +- `POST /api/analysis-runs/{run_id}/spectral-enhancements/{kind}` — on-demand spectral artifacts. `kind` is one of `cqt`, `hpss`, `onset`, `chroma_interactive`, or `reassigned` (sharper transient/frequency localization via `librosa.reassigned_spectrogram`). - `POST /api/analysis-runs/{run_id}/pitch-note-translations` - `POST /api/analysis-runs/{run_id}/interpretations` - `POST /api/analyze` (legacy compatibility) diff --git a/apps/backend/server.py b/apps/backend/server.py index 67d10859..9281154a 100644 --- a/apps/backend/server.py +++ b/apps/backend/server.py @@ -2456,6 +2456,7 @@ async def export_run_field_as_csv( "hpss": ("generate_hpss_spectrograms", ["spectrogram_harmonic", "spectrogram_percussive"], True), "onset": ("generate_onset_enhancement", ["spectrogram_onset", "onset_strength"], True), "chroma_interactive": ("generate_chroma_enhancement", ["spectrogram_chroma", "chroma_interactive"], True), + "reassigned": ("generate_reassigned_spectrogram", ["spectrogram_reassigned"], True), } diff --git a/apps/backend/spectral_viz.py b/apps/backend/spectral_viz.py index e1df6010..942e5589 100644 --- a/apps/backend/spectral_viz.py +++ b/apps/backend/spectral_viz.py @@ -426,6 +426,130 @@ def generate_onset_enhancement( return results +# Reassigned-spectrogram parameters. The magnitude floor controls how much +# of the low-energy "noise" is dropped before rendering — too aggressive +# and you lose detail; too lenient and the scatter plot turns into a wash. +# -60 dBFS below peak is a generous starting point that preserves +# attack transients and harmonic detail while hiding the floor. +REASSIGNED_MAG_FLOOR_DB = -60.0 + +# Render cap on scatter points. Each STFT cell becomes one scatter point, +# so the raw count is ``(1 + n_fft/2) * n_frames``: ~8M points on a +# 3-minute track and 15M+ on a 10-minute track at the defaults. Matplotlib's +# scatter renderer takes seconds-to-tens-of-seconds at those sizes, well +# above the librosa compute cost. 300K is a visual sweet spot — the PNG +# still looks sharp at the default ~1200×400 px figure but renders in +# well under a second. +REASSIGNED_MAX_SCATTER_POINTS = 300_000 + + +def generate_reassigned_spectrogram( + audio_path: str, + output_dir: str, + *, + sr: int = DEFAULT_SR, + n_fft: int = DEFAULT_N_FFT, + hop_length: int = DEFAULT_HOP_LENGTH, +) -> dict[str, str]: + """Generate a reassigned spectrogram PNG for sharper time/frequency localization. + + Unlike :func:`generate_spectrograms` (mel STFT on a uniform grid), + reassignment relocates each STFT bin to the local centroid of energy + in the (time, frequency) plane. Transients land at their true onset + instead of being smeared across a window; stable partials collapse + to sharp horizontal lines instead of fuzzy ridges. + + Algorithm: :func:`librosa.reassigned_spectrogram` returns three + arrays of shape ``(1 + n_fft/2, n_frames)`` — the per-bin reassigned + frequencies, per-bin reassigned times, and per-bin magnitudes. + Because the (time, freq) coordinates are no longer a uniform grid, + we render via scatter rather than :func:`librosa.display.specshow` + (which would re-rasterize to a grid and undo the sharpening). + + Returns a dict mapping artifact kind to the output file path:: + + {"spectrogram_reassigned": "/path/to/reassigned.png"} + """ + import matplotlib.figure as mpl_figure + + y = _load_audio(audio_path, sr=sr) + out = Path(output_dir) + out.mkdir(parents=True, exist_ok=True) + + # librosa returns (freqs, times, mags). Shapes are (1 + n_fft//2, + # n_frames). `fill_nan=True` substitutes the original STFT grid + # coordinate when reassignment is unstable (e.g. silent regions); + # without it, NaNs would propagate into the scatter and the plot + # would silently lose those points. + freqs, times, mags = librosa.reassigned_spectrogram( + y=y, + sr=sr, + n_fft=n_fft, + hop_length=hop_length, + reassign_frequencies=True, + reassign_times=True, + fill_nan=True, + ) + mags_db = librosa.amplitude_to_db(np.abs(mags), ref=np.max) + + # Flatten and mask to finite values above the magnitude floor. + flat_times = times.ravel() + flat_freqs = freqs.ravel() + flat_mags_db = mags_db.ravel() + mask = ( + np.isfinite(flat_times) + & np.isfinite(flat_freqs) + & np.isfinite(flat_mags_db) + & (flat_mags_db >= REASSIGNED_MAG_FLOOR_DB) + ) + times_visible = flat_times[mask] + freqs_visible = flat_freqs[mask] + mags_visible = flat_mags_db[mask] + + # Cap the number of scatter points to keep matplotlib's render time + # well under a second on long inputs. Subsampling is deterministic + # (seeded RNG) so re-running the enhancement on the same audio + # produces a byte-identical PNG — matters for idempotency tests and + # content-hash artifact caching. + if times_visible.size > REASSIGNED_MAX_SCATTER_POINTS: + rng = np.random.default_rng(0) + idx = rng.choice( + times_visible.size, + size=REASSIGNED_MAX_SCATTER_POINTS, + replace=False, + ) + times_visible = times_visible[idx] + freqs_visible = freqs_visible[idx] + mags_visible = mags_visible[idx] + + fig = mpl_figure.Figure(figsize=(FIG_WIDTH_INCHES, FIG_HEIGHT_INCHES), dpi=FIG_DPI) + ax = fig.add_axes((0, 0, 1, 1)) + ax.set_axis_off() + if times_visible.size > 0: + ax.scatter( + times_visible, + freqs_visible, + c=mags_visible, + cmap="magma", + s=0.5, + marker=",", + alpha=0.6, + linewidths=0, + vmin=REASSIGNED_MAG_FLOOR_DB, + vmax=0.0, + ) + # Axes limits: lock to the input duration and Nyquist so the PNG + # has consistent geometry regardless of how many points cleared + # the floor. + ax.set_xlim(0.0, max(1.0 / sr, len(y) / float(sr))) + ax.set_ylim(0.0, sr / 2.0) + out_path = out / "reassigned_spectrogram.png" + fig.savefig(str(out_path), dpi=FIG_DPI, bbox_inches="tight", pad_inches=0) + fig.clear() + + return {"spectrogram_reassigned": str(out_path)} + + def generate_all_artifacts( audio_path: str, output_dir: str | None = None, diff --git a/apps/backend/tests/test_spectral_viz.py b/apps/backend/tests/test_spectral_viz.py index e431aa0b..6b4c6647 100644 --- a/apps/backend/tests/test_spectral_viz.py +++ b/apps/backend/tests/test_spectral_viz.py @@ -390,5 +390,134 @@ def test_json_has_required_keys(self) -> None: self.assertEqual(len(data["timePoints"]), len(data["onsetStrength"])) +class ReassignedSpectrogramTests(unittest.TestCase): + """Tests for the opt-in librosa.reassigned_spectrogram enhancement. + + The generator is gated behind the existing + POST /api/analysis-runs/{run_id}/spectral-enhancements/reassigned + route and produces a single PNG artifact. Tests assert it runs to + completion and writes a valid PNG of non-trivial size — the visual + quality of the sharpening is empirical and not test-asserted. + """ + + def setUp(self) -> None: + self.temp_dir = tempfile.TemporaryDirectory(prefix="reassigned_test_") + self.audio_path = os.path.join(self.temp_dir.name, "test.wav") + _create_test_wav(self.audio_path) + + def tearDown(self) -> None: + self.temp_dir.cleanup() + + def test_produces_single_reassigned_png(self) -> None: + from spectral_viz import generate_reassigned_spectrogram + + out_dir = os.path.join(self.temp_dir.name, "out") + result = generate_reassigned_spectrogram(self.audio_path, out_dir) + + self.assertEqual(list(result.keys()), ["spectrogram_reassigned"]) + path = result["spectrogram_reassigned"] + self.assertTrue(os.path.isfile(path), f"reassigned PNG missing: {path}") + self.assertGreater( + os.path.getsize(path), + 1000, + "reassigned PNG suspiciously small — likely empty plot.", + ) + + def test_output_is_valid_png(self) -> None: + from spectral_viz import generate_reassigned_spectrogram + + out_dir = os.path.join(self.temp_dir.name, "out") + result = generate_reassigned_spectrogram(self.audio_path, out_dir) + with open(result["spectrogram_reassigned"], "rb") as f: + header = f.read(8) + self.assertEqual(header[:4], b"\x89PNG") + + def test_scatter_point_cap_is_enforced_for_long_inputs(self) -> None: + """Per-frame scatter rendering can produce 15M+ points on a + 10-min track. The cap (REASSIGNED_MAX_SCATTER_POINTS) keeps + matplotlib render time bounded. This test asserts the cap is + actually applied — patches ax.scatter to count call args. + """ + from unittest import mock as _mock + + from spectral_viz import ( + REASSIGNED_MAX_SCATTER_POINTS, + generate_reassigned_spectrogram, + ) + + # Generate a ~60 s audio fixture so the raw point count is well + # above the cap. At n_fft=2048 / hop=1024 / sr=44100 / 60 s + # → ~2585 frames × 1025 bins ≈ 2.6M raw points; with the floor + # mask, ~1M+ survivors — comfortably above 300K. + long_path = os.path.join(self.temp_dir.name, "long.wav") + sr = 44100 + duration = 60.0 + t = np.linspace(0, duration, int(sr * duration), endpoint=False) + # Multi-component signal so the floor mask doesn't gate too + # aggressively (a single sine would mask down to a thin line). + signal = ( + 0.4 * np.sin(2 * np.pi * 220 * t) + + 0.3 * np.sin(2 * np.pi * 660 * t) + + 0.2 * np.sin(2 * np.pi * 1320 * t) + + 0.1 * np.sin(2 * np.pi * 3000 * t) + ) + pcm = (signal * 32767).clip(-32767, 32767).astype(np.int16) + with wave.open(long_path, "w") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sr) + wf.writeframes(pcm.tobytes()) + + # Patch matplotlib's scatter to capture the call args without + # actually rendering — we only need to know how many points + # would have been drawn. + with _mock.patch( + "matplotlib.axes._axes.Axes.scatter", autospec=True + ) as mock_scatter: + out_dir = os.path.join(self.temp_dir.name, "long_out") + generate_reassigned_spectrogram(long_path, out_dir) + + self.assertTrue(mock_scatter.called, "scatter must have been called") + # First positional arg after `self` is the x array; size is the + # number of points actually rendered. + call_args = mock_scatter.call_args + x_array = call_args.args[1] if len(call_args.args) >= 2 else call_args.kwargs.get("x") + self.assertIsNotNone(x_array) + self.assertLessEqual( + len(x_array), + REASSIGNED_MAX_SCATTER_POINTS, + f"Expected at most {REASSIGNED_MAX_SCATTER_POINTS} scatter " + f"points after cap, got {len(x_array)}.", + ) + + def test_handles_silent_input_without_raising(self) -> None: + """A nearly-silent input would produce all-NaN reassignment + coordinates without ``fill_nan=True``. Guard the generator + against that by running on a silent fixture and asserting it + doesn't raise. + """ + from spectral_viz import generate_reassigned_spectrogram + + silent_path = os.path.join(self.temp_dir.name, "silent.wav") + sr = 44100 + # 2 seconds of "silence" with imperceptible dither so librosa's + # reassignment math is well-defined. + signal = np.random.default_rng(0).normal(0.0, 1e-6, sr * 2) + pcm = (signal * 32767).clip(-32767, 32767).astype(np.int16) + with wave.open(silent_path, "w") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sr) + wf.writeframes(pcm.tobytes()) + + out_dir = os.path.join(self.temp_dir.name, "silent_out") + result = generate_reassigned_spectrogram(silent_path, out_dir) + # The PNG should exist even if the scatter is sparse/empty — + # axes + colorbar geometry is still rendered. + self.assertTrue( + os.path.isfile(result["spectrogram_reassigned"]) + ) + + if __name__ == "__main__": unittest.main()