beat-detection-tcn

Beat Detection TCN

A Temporal Convolutional Network (TCN) for detecting beats in music audio. Designed for use in video editing apps where users cut video to music beats.

Quick Start

import torch
import torchaudio.transforms as T
import numpy as np
import soundfile as sf
from huggingface_hub import hf_hub_download

# Download model
model_path = hf_hub_download("finnvoorhees/beat-detection-tcn", "model.pt")
checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
config = checkpoint["config"]

# Load audio
audio, sr = sf.read("song.wav")  # or use librosa, torchaudio, etc.
if audio.ndim > 1:
    audio = audio.mean(axis=1)

# Resample to 22050 Hz if needed
if sr != config["sample_rate"]:
    import librosa
    audio = librosa.resample(audio, orig_sr=sr, target_sr=config["sample_rate"])

# Compute mel spectrogram
waveform = torch.from_numpy(audio).float().unsqueeze(0)
mel_transform = T.MelSpectrogram(
    sample_rate=config["sample_rate"], n_fft=config["n_fft"],
    hop_length=config["hop_length"], n_mels=config["n_mels"],
    f_min=config["fmin"], f_max=config["fmax"], power=2.0
)
mel = torch.log1p(mel_transform(waveform))  # (1, n_mels, n_frames)

# Load model (see train_beat_detector.py for BeatTCN class definition)
from train_beat_detector import BeatTCN, Config
cfg = Config()
for k, v in config.items():
    setattr(cfg, k, v)
model = BeatTCN(cfg)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()

# Predict
with torch.no_grad():
    logits = model(mel.unsqueeze(0) if mel.dim() == 2 else mel)
    activations = torch.sigmoid(logits).squeeze().numpy()

# Peak pick beats
fps = config["sample_rate"] / config["hop_length"]
threshold, min_interval = 0.3, 0.2
beats = []
i = 0
while i < len(activations):
    if activations[i] >= threshold:
        window_end = min(i + int(min_interval * fps), len(activations))
        peak_idx = i + np.argmax(activations[i:window_end])
        beats.append(peak_idx / fps)
        i = peak_idx + int(min_interval * fps)
    else:
        i += 1

print(f"Detected {len(beats)} beats: {beats[:10]}")

Model Details

Evaluation Results

Evaluated on synthetic test tracks spanning diverse tempos and noise conditions. F-measure with 70ms tolerance window (standard MIR evaluation metric).

BPM Estimation Accuracy

From a 120 BPM test track, the model estimated 120.3 BPM (0.25% error) via median beat interval.

Training Recipe

Based on published beat tracking research:

• All-In-One (2307.16425) — TCN architecture, 11-layer stack with exponentially growing dilations
• BeatNet (2108.03576) — Weighted BCE loss for class imbalance, 8s excerpt training

Training Data: 2,000 synthetic tracks with diverse patterns:

• 8 pattern types: click, kick, kick+snare, full drums, bass, chord stabs, layered production, sparse
• Tempo range: 60-200 BPM with humanized timing
• Background textures: pink noise, hum, pad, silence
• Various noise levels and dynamics

Training converged at epoch 25 (early stopping) with best val F-measure of 0.9999.

Usage for Video Apps

# Detect beats for video cutting
beats = detect_beats("song.mp3", threshold=0.3)

# Filter to strong beats only (for dramatic cuts)
beats_strong = detect_beats("song.mp3", threshold=0.5)

# Get beats at minimum 0.5s intervals (slower cuts)
beats_slow = detect_beats("song.mp3", min_interval=0.5)

# Estimate BPM from detected beats
intervals = np.diff(beats)
bpm = 60.0 / np.median(intervals)

Files

• model.pt — PyTorch model checkpoint (config + weights)
• train_beat_detector.py — Full training script (reproducible)
• beat_detector.py — Minimal inference script

License

MIT