finnvoorhees/beat-detection-tcn

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]}")

# 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)