From 34dd4937840710db11bf16deda600884912bb109 Mon Sep 17 00:00:00 2001 From: Zach Banks Date: Sun, 23 Oct 2022 15:40:46 -0400 Subject: [PATCH] Add Python module for testing beat tracking Hopefully @ervanalb doesn't think it's too ugly to create this new `beat_tracking_test` binary & `beat_toolkit/` folder. I had to create a mirror of the Python `rubberband` to fix some minor bugs in it. The idea is to resample, mix, and concatenate a few `.wav` files with known beat timestamps in order to generate a larger corpus of tests. --- Cargo.toml | 12 ++ beat_toolkit/Makefile | 14 ++ beat_toolkit/beat_toolkit/__init__.py | 225 ++++++++++++++++++++++++++ beat_toolkit/pyproject.toml | 44 +++++ src/bin/beat_tracking_test.rs | 29 ++++ src/lib/beat_tracking/mod.rs | 2 +- src/lib/test/frontier.beats | 42 +++++ 7 files changed, 367 insertions(+), 1 deletion(-) create mode 100644 beat_toolkit/Makefile create mode 100644 beat_toolkit/beat_toolkit/__init__.py create mode 100644 beat_toolkit/pyproject.toml create mode 100644 src/bin/beat_tracking_test.rs create mode 100644 src/lib/test/frontier.beats diff --git a/Cargo.toml b/Cargo.toml index 93689b9b..01f31d0d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + [lib] name = "radiance" path = "src/lib/mod.rs" @@ -14,6 +15,10 @@ path = "src/lib/mod.rs" name = "radiance" path = "src/bin/main.rs" +[[bin]] +name = "beat_tracking_test" +path = "src/bin/beat_tracking_test.rs" + [dependencies] env_logger = "0.9" log = "0.4" @@ -32,6 +37,13 @@ egui = "0.19" egui-wgpu = "0.19" egui-winit = "0.19" derive_more = "0.99" +# Only needed for beat_tracking_test +wav = "1.0" + +# There was a breaking change to `alsa` introduced in v0.6.1, +# temporarily pin it to v0.6.0 until this has been fixed, see: +# https://github.com/diwic/alsa-rs/issues/90 +alsa = "=0.6.0" [dev-dependencies] wav = "1.0" diff --git a/beat_toolkit/Makefile b/beat_toolkit/Makefile new file mode 100644 index 00000000..086c34fb --- /dev/null +++ b/beat_toolkit/Makefile @@ -0,0 +1,14 @@ +.PHONY: setup check run + +setup: + poetry install --all-extras + +check: + poetry run isort . + poetry run black . + poetry run flake8 . + poetry run mypy . + +run: + cd .. && cargo build --profile release --bin beat_tracking_test + poetry run pytest diff --git a/beat_toolkit/beat_toolkit/__init__.py b/beat_toolkit/beat_toolkit/__init__.py new file mode 100644 index 00000000..3217dcc5 --- /dev/null +++ b/beat_toolkit/beat_toolkit/__init__.py @@ -0,0 +1,225 @@ +import bisect +import subprocess +import sys +from dataclasses import dataclass, field, replace +from pathlib import Path +from tempfile import NamedTemporaryFile + +import numpy as np +import numpy.typing as npt +import rubberband +import soundfile + +PROFILE = "release" +BEAT_TEST_BIN = Path(__file__).parent / "../../target" / PROFILE / "beat_tracking_test" + +if not BEAT_TEST_BIN.is_file(): + print( + f"{BEAT_TEST_BIN} does not exist, please run `cargo build --profile {PROFILE}`" + ) + sys.exit(1) + + +@dataclass +class BeatSample: + time: float + is_beat: bool + activation: float + + @classmethod + def from_line(cls, line: str) -> "BeatSample": + time, is_beat, activation = line.strip().split("\t") + return cls( + time=float(time), + is_beat=is_beat == "true", + activation=float(activation), + ) + + +def get_beats(wav_path: Path) -> list[BeatSample]: + if not wav_path.is_file(): + raise ValueError(f"input file {wav_path} does not exist") + result = subprocess.run([BEAT_TEST_BIN, wav_path], check=True, capture_output=True) + lines = result.stdout.decode("utf-8").split("\n") + assert lines[0] == "t\tbeat\tactivation" + output = [] + for line in lines[1:]: + if not line: + continue + output.append(BeatSample.from_line(line)) + output = sorted(output, key=lambda x: x.time) + return output + + +def read_expected_beats(wav_path: Path) -> list[float]: + beats_path = wav_path.with_suffix(".beats") + if not beats_path.is_file(): + raise ValueError( + "beats file {beats_path} for input file {wav_path} does not exist" + ) + output = [] + for line in beats_path.read_text().split("\n"): + line = line.strip() + if not line: + continue + output.append(float(line)) + output = sorted(output) + return output + + +@dataclass +class BeatComparison: + expected: list[float] = field(repr=False) + measured: list[BeatSample] = field(repr=False) + + exact_match: bool = field() + error_sq: float = field() + aligned_fraction: float = field() + + @classmethod + def compare( + cls, + expected: list[float], + measured: list[BeatSample], + sample_rate: float = 0.01, + ) -> "BeatComparison": + measured_beat_times = [b.time for b in measured if b.is_beat] + + # This measurement is asymmetric, crude, and slow to compute, + # but it's good enough for some rough tests/comparisons + error_sq = 0.0 + exact_count = 0 + for exp in expected: + left = bisect.bisect_left(measured_beat_times, exp) + right = bisect.bisect_right(measured_beat_times, exp) + error = min( + abs(exp - m) for m in measured_beat_times[max(0, left - 1) : right + 1] + ) + # Allow rounding from unaligned sample rates + error = max(0, error - sample_rate / 2) + if error == 0: + exact_count += 1 + error_sq += error**2 + + exact_match = (exact_count == len(expected)) and len(expected) == len( + measured_beat_times + ) + + return cls( + expected=expected, + measured=measured, + exact_match=exact_match, + error_sq=error_sq / len(expected), + aligned_fraction=exact_count / len(expected), + ) + + +@dataclass +class Wav: + name: str = field() + sample_rate: int = field() + data: npt.NDArray[np.int16] = field(repr=False) + expected_beats: list[float] = field(repr=False) + + def __post_init__(self) -> None: + assert sorted(self.expected_beats) == self.expected_beats + assert len(self.expected_beats) >= 2 + assert 0.0 <= self.expected_beats[0] <= 3.0 + assert 0.0 <= self.duration_seconds() - self.expected_beats[-1] <= 3.0 + + @classmethod + def from_path(cls, wav_path: Path) -> "Wav": + name = wav_path.name + data, sample_rate = soundfile.read(wav_path, dtype="int16") + expected_beats = read_expected_beats(wav_path) + + return cls( + name=name, + sample_rate=sample_rate, + data=data, + expected_beats=expected_beats, + ) + + def duration_seconds(self) -> float: + return len(self.data) / self.sample_rate + + def get_beats(self) -> list[BeatSample]: + with NamedTemporaryFile(suffix=".wav") as tmp: + soundfile.write(tmp, self.data, self.sample_rate) + tmp.flush() + return get_beats(Path(tmp.name)) + + def stretch(self, ratio: float) -> "Wav": + new_data = rubberband.stretch( + self.data, + rate=self.sample_rate, + ratio=ratio, + crispness=5, + formants=False, + precise=True, + ) + new_beats = [t * ratio for t in self.expected_beats] + return replace( + self, + data=new_data, + expected_beats=new_beats, + ) + + def compare(self) -> BeatComparison: + return BeatComparison.compare(self.expected_beats, self.get_beats()) + + def scale(self, k: float) -> "Wav": + # TODO: saturating math? + return replace( + self, + data=self.data * k, + ) + + @classmethod + def concat(cls, *wavs: "Wav") -> "Wav": + if not wavs: + raise ValueError("must provide at least one Wav") + sample_rate = wavs[0].sample_rate + if not all(wav.sample_rate == sample_rate for wav in wavs): + raise ValueError("all wavs must have the same sample rate") + name = "|".join(wav.name for wav in wavs) + data = np.concatenate([wav.data for wav in wavs]) + + expected_beats: list[float] = [] + t = 0.0 + for wav in wavs: + expected_beats.extend(b + t for b in wav.expected_beats) + t += wav.duration_seconds() + + return cls( + name=name, + sample_rate=sample_rate, + data=data, + expected_beats=expected_beats, + ) + + +def test_frontier_wav() -> None: + example_path = BEAT_TEST_BIN.parent / "../../src/lib/test/frontier.wav" + wav = Wav.from_path(example_path) + measured = wav.get_beats() + for b in measured: + if b.is_beat: + print(b) + comparison = wav.compare() + print(comparison) + assert comparison.exact_match + + for ratio in [0.5, 0.7, 0.9, 1.1, 1.3, 1.5, 1.9]: + print(ratio, wav.stretch(ratio).compare()) + + concat = Wav.concat(wav, wav, wav, wav) + for b in concat.get_beats(): + if b.is_beat: + print(b) + print(concat.expected_beats) + print("concat", concat.compare()) + + +if __name__ == "__main__": + test_frontier_wav() diff --git a/beat_toolkit/pyproject.toml b/beat_toolkit/pyproject.toml new file mode 100644 index 00000000..f53700d2 --- /dev/null +++ b/beat_toolkit/pyproject.toml @@ -0,0 +1,44 @@ +[tool.poetry] +name = "beat-toolkit" +version = "0.1.0" +description = "Beat tracking testing toolkit" +authors = ["Zach Banks "] +packages = [{include = "beat_toolkit"}] + +[tool.poetry.dependencies] +python = "^3.10" +numpy = "^1.23.4" +soundfile = "^0.11.0" +rubberband = {git = "https://github.com/zbanks/rubberband.git", rev = "v1.0.3"} + + +[tool.poetry.group.dev.dependencies] +black = "^22.10.0" +flake8 = "^5.0.4" +mypy = "^0.982" +isort = "^5.10.1" +pytest = "^7.1.3" +Flake8-pyproject = "^1.1.0.post0" + +[build-system] +requires = ["poetry-core"] +build-backend = "poetry.core.masonry.api" + +[tool.mypy] +plugins = "numpy.typing.mypy_plugin" +strict = true + +[[tool.mypy.overrides]] +module = [ + "rubberband", + "soundfile", +] +ignore_missing_imports = true + +[tool.isort] +profile = "black" + +[tool.flake8] +max-line-length = 88 +exclude = ".git,.mypy_cache,.pytest_cache" +extend-ignore = "E203" diff --git a/src/bin/beat_tracking_test.rs b/src/bin/beat_tracking_test.rs new file mode 100644 index 00000000..ef4174a5 --- /dev/null +++ b/src/bin/beat_tracking_test.rs @@ -0,0 +1,29 @@ +use radiance::{BeatTracker, FPS}; +use std::fs::File; +use std::path::Path; + +fn main() { + // Usage: ./beat_tracking_test + let mut args = std::env::args(); + args.next().unwrap(); + let filename = args.next().expect("expect .wav file"); + + // Read music from audio file + let mut inp_file = File::open(Path::new(&filename)).unwrap(); + let (header, data) = wav::read(&mut inp_file).unwrap(); + assert_eq!(header.audio_format, wav::WAV_FORMAT_PCM); + assert_eq!(header.channel_count, 1); + assert_eq!(header.sampling_rate, 44100); + assert_eq!(header.bits_per_sample, 16); + let data = data.try_into_sixteen().unwrap(); + + // Instantiate a BeatTracker + let mut bt = BeatTracker::new(); + let beats = bt.process(&data); + + // Print results as a TSV + println!("{}\t{}\t{}", "t", "beat", "activation"); + for (i, &(_, activation, beat)) in beats.iter().enumerate() { + println!("{}\t{}\t{}", (i as f32) / FPS, beat, activation); + } +} diff --git a/src/lib/beat_tracking/mod.rs b/src/lib/beat_tracking/mod.rs index b818d367..ec42f6cd 100644 --- a/src/lib/beat_tracking/mod.rs +++ b/src/lib/beat_tracking/mod.rs @@ -43,7 +43,7 @@ use nalgebra::{DVector, DMatrix, SVector}; use itertools::iproduct; pub const SAMPLE_RATE: f32 = 44100.; -const FPS: f32 = 100.; +pub const FPS: f32 = 100.; // Hop size of 441 with sample rate of 44100 Hz gives an output frame rate of 100 Hz const HOP_SIZE: usize = 441; const FRAME_SIZE: usize = 2048; diff --git a/src/lib/test/frontier.beats b/src/lib/test/frontier.beats new file mode 100644 index 00000000..34b6b320 --- /dev/null +++ b/src/lib/test/frontier.beats @@ -0,0 +1,42 @@ +0.51 +1.47 +2.21 +2.93 +3.66 +4.39 +4.88 +5.35 +5.83 +6.31 +6.79 +7.27 +7.77 +8.25 +8.73 +9.21 +9.70 +10.20 +10.67 +11.15 +11.63 +12.11 +12.59 +13.07 +13.57 +14.10 +14.58 +15.08 +15.51 +15.99 +16.47 +16.95 +17.44 +17.94 +18.43 +18.92 +19.38 +19.85 +20.35 +20.84 +21.32 +21.80