Beat This! Beat Tracker

repository·main·Indexed 43 days ago

https://github.com/cpjku/beat_this

An accurate beat tracking system presented at ISMIR 2024 that predicts beats and downbeats from audio files using transformer-based models. It provides a Python API and a command-line interface for inference, as well as tools for training models and computing evaluation metrics. Key features include the ShiftTolerantBCELoss, support for torch.hub loading, and integration with madmom's DBN post-processing.

Tokens
5.9K
Snippets
20
Records
26
Agent score
57%

What's inside Beat This!

  1. Install requirements for reproducing paper metrics

    main

    To compute evaluation metrics as presented in the paper, you need to install pytorch_lightning, pandas, and a specific version of mir_eval from source to ensure compatibility with numpy versions.

    Note: You must have the annotations and spectrogram datasets set up. The GTZAN dataset is sufficient for --datasplit test, but all other datasets are required for --datasplit val.

    pip install pytorch_lightning pandas
    pip install https://github.com/mir-evaluation/mir_eval/archive/main.zip
  2. Set up training data and annotations

    main

    To train or evaluate the models, you need annotations and spectrograms.

    1. Download Annotations

    Annotations are hosted in a separate repository. Clone them to data/annotations:

    mkdir -p data
    git clone https://github.com/CPJKU/beat_this_annotations data/annotations
    # Optional: Use version 1.0 for exact paper reproducibility
    # cd data/annotations; git checkout v1.0

    2. Download Spectrograms

    Spectrograms are available as a Zenodo dataset.

    • For GTZAN test set: Download gtzan.zip.
    • For Validation/Training: Download all available .zip files (except beat_this_annotations.zip).
    • Placement: Extract all .zip files into data/audio/spectrograms/. The code supports .npz files (e.g., data/audio/spectrograms/gtzan.npz) or directories of .npy files.

    3. Recreating Spectrograms from Audio

    If you have original audio files, you can generate your own spectrograms:

    1. Create data/audio_paths.tsv with the format: [dataset_name]<TAB>[path_to_audio_directory].
    2. Ensure corresponding annotations are in data/annotations.
    3. Install requirements: pip install pandas pedalboard.
    4. Run the preprocessing script:
    python launch_scripts/preprocess_audio.py
  3. Manually integrate the BeatThis model

    main

    To copy the model into a custom project without using the package, you must include:

    1. beat_this/model/beat_tracker.py
    2. beat_this/model/roformer.py

    Dependencies: einops, rotary-embedding-torch.

    Note: If you remove the BeatThis.state_dict() and BeatThis._load_from_state_dict() methods (which are workarounds for compiled models), there are no other internal dependencies.

  4. Install Beat This!

    main

    To use the beat tracker, you need to install PyTorch, several dependencies, and the beat-this package itself.

    1. Install PyTorch

    Install PyTorch 2.0 or later according to your platform's instructions.

    2. Install Dependencies

    Install the required Python modules using pip:

    pip install tqdm einops soxr rotary-embedding-torch

    Note: If using conda, rotary-embedding-torch is only available on PyPI.

    3. Install Audio Backends

    To support audio formats other than .wav, ensure ffmpeg or another supported torchaudio backend is installed on your system.

    4. Install the Package

    Install the stable version:

    pip install beat-this

    Or the development version from GitHub:

    pip install https://github.com/CPJKU/beat_this/archive/main.zip
    pip install beat-this
  5. Load the BeatThis model via torch.hub

    main

    You can use the model without installing the package by loading it via torch.hub.

    import torch
    beat_this = torch.hub.load('CPJKU/beat_this', 'beat_this', 'final0', device='cuda')
  6. Load the BeatThis model via Python package

    main

    If you have installed the beat_this package, you can import the model class directly or use the load_model utility to load pretrained weights.

    • BeatThis: An untrained model that maps spectrograms to frame-wise beat and downbeat logits.
    • load_model(model_name, device): Loads a pretrained model (e.g., 'final0') to a specific device.
    # Import the class for an untrained model
    from beat_this.model.beat_tracker import BeatThis
    
    # Load a pretrained model
    from beat_this.inference import load_model
    beat_this = load_model('final0', device='cuda')
  7. Use the File2Beats Python class

    main

    For programmatic access, use the beat_this.inference module. The File2Beats class handles model loading, audio pre-processing, and beat prediction.

    Inference with Files

    from beat_this.inference import File2Beats
    
    # Initialize the tracker
    # checkpoint_path: name of the model (e.g., 'final0')
    # device: 'cuda' or 'cpu'
    # dbn: boolean, whether to use DBN postprocessing
    file2beats = File2Beats(checkpoint_path="final0", device="cuda", dbn=False)
    
    # Predict beats and downbeats
    audio_path = "path/to/audio.file"
    beats, downbeats = file2beats(audio_path)

    Saving Results

    You can save the predicted beats to a .beats file (compatible with Sonic Visualizer) using save_beat_tsv:

    from beat_this.utils import save_beat_tsv
    
    outpath = "path/to/output.beats"
    save_beat_tsv(beats, downbeats, outpath)

    Other Inference Classes

    • Audio2Beats: Use this if you already have an audio tensor loaded (requires passing the tensor and sample rate).
    • Audio2Frames: Provides framewise logits.
    • Spect2Frames: Provides spectrogram inputs.
    from beat_this.inference import File2Beats
    file2beats = File2Beats(checkpoint_path="final0", device="cuda", dbn=False)
    beats, downbeats = file2beats("path/to/audio.file")
  8. Transform waveforms to Log-Mel Spectrograms with LogMelSpect

    main

    The LogMelSpect class is a torch.nn.Module that transforms a 1D waveform into a 2D log-mel spectrogram.

    Input: A waveform as a monodimensional array/tensor of shape (T). Output: A 2D log mel spectrogram of shape (F, 128), where F is the number of mel bins.

    Initialization Parameters:

    • sample_rate (int): The sampling rate of the input audio. Default is 22050.
    • n_fft (int): The size of the FFT. Default is 1024.
    • hop_length (int): The distance between adjacent frames. Default is 441.
    • f_min (int): Minimum frequency in Hz. Default is 30.
    • f_max (int): Maximum frequency in Hz. Default is 11000.
    • n_mels (int): Number of mel filter banks. Default is 128.
    • mel_scale (str): The mel scale to use. Default is `
  9. Select available models

    main

    Models are automatically downloaded during inference if not present. You can specify a model using the --model CLI flag or the checkpoint_path parameter in Python.

    Main Models (~78 MB each)

    • final0, final1, final2: Main models trained on all data except GTZAN (using three different seeds).
    • small0, small1, small2: Smaller versions (~8.1 MB each).
    • single_final0, single_final1, single_final2: Main models trained on a single split.
    • fold0 through fold7: Models trained in an 8-fold cross-validation setting.

    Other Models

    • hung0, hung1, hung2: Models trained on data used by the Hung et al. system.
    • Various ablation models (e.g., single_notempoaug0, single_nosumhead0) are available for research reproducibility.

    Note on Evaluation: To avoid biased results, do not evaluate using final* or small* models on training datasets. For fair evaluation on datasets used in the paper, use single_final* (validation part of single split) or fold* (8-fold cross-validation).

  10. Use LogMelSpect for spectrogram transformation

    main

    To use LogMelSpect, initialize the module with your desired audio parameters and pass a waveform tensor to the forward method (or call the instance directly).

    import torch
    from beat_this.preprocessing import LogMelSpect
    
    # Initialize the transform
    transform = LogMelSpect(sample_rate=22050, n_mels=128)
    
    # Create a dummy waveform tensor of shape (T)
    waveform = torch.randn(22050)
    
    # Compute the log-mel spectrogram
    spectrogram = transform(waveform)
    # Output shape will be (128, T_frames)
  11. Compute paper metrics via command line

    main

    Use launch_scripts/compute_paper_metrics.py to evaluate models on different datasets and configurations.

    Key arguments:

    • --models: A space-separated list of model names (e.g., final0 final1 final2).
    • --datasplit: Use test for GTZAN or val for other datasets.
    • --aggregation-type: Use k-fold for cross-validation results.
    • --dbn: Enables DBN postprocessing (requires the madmom package).
    # Compute main results on test set
    python launch_scripts/compute_paper_metrics.py --models final0 final1 final2 --datasplit test
    
    # Compute 8-fold cross-validation (Table 1)
    python launch_scripts/compute_paper_metrics.py --models fold0 fold1 fold2 fold3 fold4 fold5 fold6 fold7 --datasplit val --aggregation-type k-fold
    
    # Compute with DBN postprocessing
    python launch_scripts/compute_paper_metrics.py --models final0 final1 final2 --datasplit test --dbn