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.
1import torch
2import torchaudio.transforms as T
3import numpy as np
4import soundfile as sf
5from huggingface_hub import hf_hub_download
6
7# Download model
8model_path = hf_hub_download("finnvoorhees/beat-detection-tcn", "model.pt")
9checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
10config = checkpoint["config"]
11
12# Load audio
13audio, sr = sf.read("song.wav") # or use librosa, torchaudio, etc.
14if audio.ndim > 1:
15 audio = audio.mean(axis=1)
16
17# Resample to 22050 Hz if needed
18if sr != config["sample_rate"]:
19 import librosa
20 audio = librosa.resample(audio, orig_sr=sr, target_sr=config["sample_rate"])
21
22# Compute mel spectrogram
23waveform = torch.from_numpy(audio).float().unsqueeze(0)
24mel_transform = T.MelSpectrogram(
25 sample_rate=config["sample_rate"], n_fft=config["n_fft"],
26 hop_length=config["hop_length"], n_mels=config["n_mels"],
27 f_min=config["fmin"], f_max=config["fmax"], power=2.0
28)
29mel = torch.log1p(mel_transform(waveform)) # (1, n_mels, n_frames)
30
31# Load model (see train_beat_detector.py for BeatTCN class definition)
32from train_beat_detector import BeatTCN, Config
33cfg = Config()
34for k, v in config.items():
35 setattr(cfg, k, v)
36model = BeatTCN(cfg)
37model.load_state_dict(checkpoint["model_state_dict"])
38model.eval()
39
40# Predict
41with torch.no_grad():
42 logits = model(mel.unsqueeze(0) if mel.dim() == 2 else mel)
43 activations = torch.sigmoid(logits).squeeze().numpy()
44
45# Peak pick beats
46fps = config["sample_rate"] / config["hop_length"]
47threshold, min_interval = 0.3, 0.2
48beats = []
49i = 0
50while i < len(activations):
51 if activations[i] >= threshold:
52 window_end = min(i + int(min_interval * fps), len(activations))
53 peak_idx = i + np.argmax(activations[i:window_end])
54 beats.append(peak_idx / fps)
55 i = peak_idx + int(min_interval * fps)
56 else:
57 i += 1
58
59print(f"Detected {len(beats)} beats: {beats[:10]}")
Evaluated on synthetic test tracks spanning diverse tempos and noise conditions. F-measure with 70ms tolerance window (standard MIR evaluation metric).
1# Detect beats for video cutting
2beats = detect_beats("song.mp3", threshold=0.3)
3
4# Filter to strong beats only (for dramatic cuts)
5beats_strong = detect_beats("song.mp3", threshold=0.5)
6
7# Get beats at minimum 0.5s intervals (slower cuts)
8beats_slow = detect_beats("song.mp3", min_interval=0.5)
9
10# Estimate BPM from detected beats
11intervals = np.diff(beats)
12bpm = 60.0 / np.median(intervals)
