Music onset detector from Dance Dance Convolution packaged as a lightweight PyTorch module
from .paths import LIB_DIR
from .constants import SAMPLE_RATE, FRAME_RATE, Difficulty
from .spectral import SpectrogramExtractor
from .cnn import SpectrogramNormalizer, PlacementCNN
from .util import compute_onset_salience, find_peaks, threshold_peaks


# NOTE: This changes the test discovery pattern from "test*.py" (default) to "*test.py".
def load_tests(loader, standard_tests, pattern):
    package_tests = loader.discover(start_dir=LIB_DIR, pattern="*test.py")
    standard_tests.addTests(package_tests)
    return standard_tests

ddc_onset/__init__.py
chrisdonahue/ddc_onset
import pathlib

import torch
import torch.nn as nn
import torch.nn.functional as F

from .constants import Difficulty
from .paths import WEIGHTS_DIR

_NUM_MEL_BANDS = 80
_NUM_FFT_FRAME_LENGTHS = 3


class SpectrogramNormalizer(nn.Module):
    """Normalizes log-Mel spectrograms to zero mean and unit variance per bin."""

    def __init__(self, load_moments: bool = True):
        super().__init__()
        self.mean = nn.Parameter(
            torch.zeros(
                (_NUM_MEL_BANDS, _NUM_FFT_FRAME_LENGTHS),
                dtype=torch.float32,
                requires_grad=False,
            ),
            requires_grad=False,
        )
        self.std = nn.Parameter(
            torch.ones(
                (_NUM_MEL_BANDS, _NUM_FFT_FRAME_LENGTHS),
                dtype=torch.float32,
                requires_grad=False,
            ),
            requires_grad=False,
        )
        if load_moments:
            self.load_state_dict(
                torch.load(pathlib.Path(WEIGHTS_DIR, "spectrogram_normalizer.bin"))
            )

    def forward(self, x: torch.Tensor):
        """Normalizes log-Mel spectrograms to zero mean and unit variance per bin.

        Args:
            x: 44.1kHz waveforms as float32 [batch_size, num_frames, num_mel_bands (80), num_fft_frame_lengths (3)].
        Returns:
            Normalized input (same shape).
        """
        return (x - self.mean) / self.std


_FEATURE_CONTEXT_RADIUS_1 = 7
_FEATURE_CONTEXT_RADIUS_2 = 3


class PlacementCNN(nn.Module):
    """Predicts placement scores from log-Mel spectrograms."""

    def __init__(self, load_pretrained_weights: bool = True):
        super().__init__()
        self.conv0 = nn.Conv2d(3, 10, (7, 3))
        self.maxpool0 = nn.MaxPool2d((1, 3), (1, 3))
        self.conv1 = nn.Conv2d(10, 20, (3, 3))
        self.maxpool1 = nn.MaxPool2d((1, 3), (1, 3))
        self.dense0 = nn.Linear(1125, 256)
        self.dense1 = nn.Linear(256, 128)
        self.output = nn.Linear(128, 1)
        if load_pretrained_weights:
            self.load_state_dict(
                torch.load(pathlib.Path(WEIGHTS_DIR, "placement_cnn_ckpt_56000.bin"))
            )

    def conv(self, x: torch.Tensor):
        # x is b, 3, 15, 80

        # Conv 0
        x = self.conv0(x)
        x = F.relu(x)
        x = self.maxpool0(x)

        # Conv 1
        x = self.conv1(x)
        x = F.relu(x)
        x = self.maxpool1(x)

        return x

    def dense(self, x_conv_diff: torch.Tensor):
        # x is b, 1125

        # Dense 0
        x = self.dense0(x_conv_diff)
        x = F.relu(x)

        # Dense 1
        x = self.dense1(x)
        x = F.relu(x)

        # Output
        x = self.output(x)
        x = x.view(-1)

        return x

    def forward(
        self,
        x: torch.Tensor,
        difficulties: torch.Tensor,
        output_logits: bool = False,
        conv_chunk_size: int = 256,
        dense_chunk_size: int = 256,
    ):
        """Predicts placement scores from normalized log-Mel spectrograms.

        Args:
            x: Normalized Log-Mel spectrograms as float32 [num_frames, num_mel_bands (80), num_fft_frame_lengths (3)].
            difficulties: DDR difficulty labels as int64 [batch_size]
            output_logits: If True, output raw logits instead of sigmoid scores (default).

        Returns:
            Placement scores (or logits) as float32 [batch_size, num_frames].
        """

        # TODO: Proper batch support for this module

        # x is t, 80, 3
        num_timesteps = x.shape[0]

        # Pad features
        x_padded = F.pad(
            x, (0, 0, 0, 0, _FEATURE_CONTEXT_RADIUS_1, _FEATURE_CONTEXT_RADIUS_1)
        )

        # Convolve
        x_padded = x_padded.permute(2, 0, 1)
        x_conv = []
        for i in range(0, num_timesteps, conv_chunk_size):
            x_chunk = x_padded[
                :, i : i + conv_chunk_size + _FEATURE_CONTEXT_RADIUS_1 * 2
            ].unsqueeze(0)
            x_chunk_conv = self.conv(x_chunk)
            assert x_chunk_conv.shape[1] > _FEATURE_CONTEXT_RADIUS_2 * 2
            if i == 0:
                x_conv.append(x_chunk_conv[:, :, :_FEATURE_CONTEXT_RADIUS_2])
            x_conv.append(
                x_chunk_conv[:, :, _FEATURE_CONTEXT_RADIUS_2:-_FEATURE_CONTEXT_RADIUS_2]
            )
        x_conv.append(x_chunk_conv[:, :, -_FEATURE_CONTEXT_RADIUS_2:])
        x_conv = torch.cat(x_conv, dim=2)
        x_conv = x_conv.permute(0, 2, 3, 1)
        x_conv = x_conv.reshape(-1, 160)

        # Dense
        logits = []
        for i in range(0, num_timesteps, dense_chunk_size):
            # TODO: Turn this into a convolutional layer?
            # NOTE: Pytorch didn't like this as of 20-03-15:
            # https://github.com/pytorch/pytorch/pull/33073
            x_chunk = []
            for j in range(i, i + dense_chunk_size):
                if j >= num_timesteps:
                    break
                x_chunk.append(x_conv[j : j + 1 + _FEATURE_CONTEXT_RADIUS_2 * 2])
            x_chunk = torch.stack(x_chunk, dim=0)
            x_chunk = x_chunk.reshape(-1, 1120)

            # Compute dense layer for each difficulty
            logits_diffs = []
            for k in range(difficulties.shape[0]):
                d = difficulties[k].repeat(x_chunk.shape[0])
                doh = F.one_hot(d, len(Difficulty)).float()
                x_chunk_diff = torch.cat([x_chunk, doh], dim=1)
                x_chunk_dense = self.dense(x_chunk_diff)
                logits_diffs.append(x_chunk_dense)
            logits_diffs = torch.stack(logits_diffs, dim=0)
            logits.append(logits_diffs)
        logits = torch.cat(logits, dim=1)

        if output_logits:
            return logits
        else:
            return torch.sigmoid(logits)

ddc_onset/cnn.py
chrisdonahue/ddc_onset
import os
import pickle
import unittest

import numpy as np
import torch

from .constants import Difficulty, DIFFICULTY_TO_THRESHOLD
from .paths import TEST_DATA_DIR
from .cnn import SpectrogramNormalizer, PlacementCNN
from .util import find_peaks, threshold_peaks

_SHORT_FEATS_REF_FP = os.path.join(TEST_DATA_DIR, "short_essentia_feats_ref.npy")
_SHORT_SCORES_REF_FP = os.path.join(
    TEST_DATA_DIR, "short_essentia_feats_ref_ddc_placement_scores.pkl"
)
_LONG_FEATS_REF_FP = os.path.join(TEST_DATA_DIR, "long_essentia_feats_ref.npy")
_LONG_SCORES_REF_FP = os.path.join(
    TEST_DATA_DIR, "long_essentia_feats_ref_ddc_placement_scores.pkl"
)

_SHORT_DIFF_TO_NUM_PEAKS = {0: (69, 0), 1: (69, 0), 2: (68, 1), 3: (68, 3), 4: (66, 7)}

_LONG_DIFF_TO_NUM_PEAKS = {
    0: (3051, 718),
    1: (3033, 815),
    2: (3025, 942),
    3: (3027, 1170),
    4: (3031, 1288),
}


class TestPlacement(unittest.TestCase):
    def setUp(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.normalizer = SpectrogramNormalizer()
        self.normalizer.to(self.device)

        self.placement = PlacementCNN()
        self.placement.to(self.device)

    def test_placement(self):
        for feats_fp, scores_ref_fp, diff_to_num_peaks in zip(
            [_SHORT_FEATS_REF_FP, _LONG_FEATS_REF_FP],
            [_SHORT_SCORES_REF_FP, _LONG_SCORES_REF_FP],
            [_SHORT_DIFF_TO_NUM_PEAKS, _LONG_DIFF_TO_NUM_PEAKS],
        ):
            feats = np.load(feats_fp)
            feats = torch.tensor(feats, device=self.device, requires_grad=False)
            feats_norm = self.normalizer(feats)

            with open(scores_ref_fp, "rb") as f:
                diff_to_scores_ref = pickle.load(f, encoding="latin1")
            scores_ref = [diff_to_scores_ref[d.value] for d in Difficulty]
            scores_ref = np.array(scores_ref)

            for conv_chunk_size in [128, 256, 512]:
                for dense_chunk_size in [128, 256, 512]:
                    with torch.no_grad():
                        diffs = torch.tensor(
                            [d.value for d in Difficulty],
                            device=self.device,
                            dtype=torch.int64,
                        )
                        scores = self.placement.forward(
                            feats_norm,
                            diffs,
                            conv_chunk_size=conv_chunk_size,
                            dense_chunk_size=dense_chunk_size,
                        )
                        scores = scores.cpu().numpy()

                    error = np.sum(np.abs(scores - scores_ref))
                    self.assertTrue(error < 2e-3)

                    for d in Difficulty:
                        peaks = find_peaks(scores[d.value])
                        thresholded_peaks = threshold_peaks(
                            scores[d.value],
                            peaks,
                            threshold=DIFFICULTY_TO_THRESHOLD[d],
                        )
                        self.assertEqual(
                            (len(peaks), len(thresholded_peaks)),
                            diff_to_num_peaks[d.value],
                        )


if __name__ == "__main__":
    unittest.main()

ddc_onset/cnn_test.py
chrisdonahue/ddc_onset
from enum import Enum

SAMPLE_RATE = 44100
FRAME_RATE = 100


# Coarse difficulties from DDR. For general onset detection use, CHALLENGE is a good default.
class Difficulty(Enum):
    BEGINNER = 0
    EASY = 1
    MEDIUM = 2
    HARD = 3
    CHALLENGE = 4


# Thresholds tuned on DDR validation data. Here for testing / posterity and likely irrelevant for other applications.
DIFFICULTY_TO_THRESHOLD = {
    Difficulty.BEGINNER: 0.15325437,
    Difficulty.EASY: 0.23268291,
    Difficulty.MEDIUM: 0.29456162,
    Difficulty.HARD: 0.29084727,
    Difficulty.CHALLENGE: 0.28875697,
}

ddc_onset/constants.py
chrisdonahue/ddc_onset
import pathlib

LIB_DIR = pathlib.Path(__file__).resolve().parent
_REPO_DIR = LIB_DIR.parent if LIB_DIR.parent.name == "ddc_onset" else None

WEIGHTS_DIR = pathlib.Path(LIB_DIR, "weights")
TEST_DATA_DIR = None if _REPO_DIR is None else pathlib.Path(_REPO_DIR, "test")

ddc_onset/paths.py
chrisdonahue/ddc_onset
from typing import Optional
import pathlib

import torch
import torch.nn as nn
import torch.nn.functional as F

from .constants import SAMPLE_RATE, FRAME_RATE
from .paths import WEIGHTS_DIR

_FEATS_HOP = SAMPLE_RATE // FRAME_RATE
_FFT_FRAME_LENGTHS = [1024, 2048, 4096]
_NUM_MEL_BANDS = 80
_LOG_EPS = 1e-16


class SpectrogramExtractor(nn.Module):
    """Extracts log-Mel spectrograms from waveforms."""

    def __init__(self):
        super().__init__()
        for i in _FFT_FRAME_LENGTHS:
            window = nn.Parameter(
                torch.zeros(i, dtype=torch.float32, requires_grad=False),
                requires_grad=False,
            )
            setattr(self, "window_{}".format(i), window)
            mel = nn.Parameter(
                torch.zeros(
                    ((i // 2) + 1, 80), dtype=torch.float32, requires_grad=False
                ),
                requires_grad=False,
            )
            setattr(self, "mel_{}".format(i), mel)
        self.load_state_dict(
            torch.load(pathlib.Path(WEIGHTS_DIR, "spectrogram_extractor.bin"))
        )

    def forward(
        self, x: torch.Tensor, frame_chunk_size: Optional[int] = None
    ) -> torch.Tensor:
        """Extracts an 80-bin log-Mel spectrogram from a waveform with 3 different FFT frame lengths.

        Args:
            x: 44.1kHz waveforms as float32 [batch_size, num_samples].
            frame_chunk_size: Number of frames to process at a time. If None, process all frames at once.
        Returns:
            Log mel spectrograms as float32 [batch_size, num_frames, num_mel_bands (80), num_fft_frame_lengths (3)].
        """
        # NOTE: This was originally implemented as [samps, batch] but [batch, samps] is better type signature.
        waveform = x.transpose(1, 0)
        feats = []
        for i, fft_frame_length in enumerate(_FFT_FRAME_LENGTHS):
            # Pad waveform to center spectrogram
            fft_frame_length_half = fft_frame_length // 2
            waveform_padded = F.pad(waveform, (0, 0, fft_frame_length_half, 0))
            waveform_padded_len = waveform_padded.shape[0]

            # Params for this FFT size
            window = getattr(self, "window_{}".format(fft_frame_length)).view(1, -1, 1)
            mel_w = getattr(self, "mel_{}".format(fft_frame_length))

            # Chunk up waveform to save memory at cost of some efficiency
            if frame_chunk_size is None:
                chunk_hop = waveform_padded_len
            else:
                chunk_hop = frame_chunk_size * _FEATS_HOP

            chunk_feats = []
            for c in range(0, waveform_padded_len, chunk_hop):
                frames = []
                for s in range(c, min(c + chunk_hop, waveform_padded_len), _FEATS_HOP):
                    # Slice waveform into frames
                    # TODO: Change this to range(0, waveform.shape[0], _FEATS_HOP) to make num feats = ceil(num_samps / 441)? Make sure frames are equal after doing this
                    frame = waveform_padded[s : s + fft_frame_length]
                    padding_amt = fft_frame_length - frame.shape[0]
                    if padding_amt > 0:
                        frame = F.pad(frame, (0, 0, padding_amt, 0))
                    frames.append(frame)
                frames = torch.stack(frames, dim=0)

                # Apply window
                frames *= window

                # Copying questionable "zero phase" windowing nonsense from essentia
                # https://github.com/MTG/essentia/blob/master/src/algorithms/standard/windowing.cpp#L85
                frames_half_one = frames[:, :fft_frame_length_half]
                frames_half_two = frames[:, fft_frame_length_half:]
                frames = torch.cat([frames_half_two, frames_half_one], dim=1)

                # Perform FFT
                frames = torch.transpose(frames, 1, 2)
                spec = torch.fft.rfft(frames)

                # Compute power spectrogram
                spec_r = torch.real(spec)
                spec_i = torch.imag(spec)
                pow_spec = torch.pow(spec_r, 2) + torch.pow(spec_i, 2)

                # Compute mel spectrogram
                mel_spec = torch.matmul(pow_spec, mel_w)

                # Compute log mel spectrogram
                log_mel_spec = torch.log(mel_spec + _LOG_EPS)
                log_mel_spec = torch.transpose(log_mel_spec, 1, 2)
                chunk_feats.append(log_mel_spec)

            chunk_feats = torch.cat(chunk_feats, dim=0)
            # Cut off extra chunk_feats for larger FFT lengths
            # TODO: Don't slice these chunk_feats to begin with
            if i == 0:
                feats_num_timesteps = chunk_feats.shape[0]
            else:
                assert chunk_feats.shape[0] > feats_num_timesteps
                chunk_feats = chunk_feats[:feats_num_timesteps]

            feats.append(chunk_feats)

        feats = torch.stack(feats, dim=2)

        # [feats..., batch] -> [batch, feats...]
        return feats.permute(3, 0, 1, 2)

ddc_onset/spectral.py
chrisdonahue/ddc_onset
import pathlib
import unittest

import numpy as np
import torch

from .paths import TEST_DATA_DIR
from .spectral import SpectrogramExtractor

_SHORT_AUDIO_REF_FP = pathlib.Path(TEST_DATA_DIR, "short_essentia_audio_ref.npy")
_SHORT_FEATS_REF_FP = pathlib.Path(TEST_DATA_DIR, "short_essentia_feats_ref.npy")
_LONG_AUDIO_REF_FP = pathlib.Path(TEST_DATA_DIR, "long_essentia_audio_ref.npy")
_LONG_FEATS_REF_FP = pathlib.Path(TEST_DATA_DIR, "long_essentia_feats_ref.npy")


class TestFeatureExtraction(unittest.TestCase):
    def setUp(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.extractor = SpectrogramExtractor()
        self.extractor.to(self.device)

    def test_extract_feats(self):
        for batch_size in [1, 2]:
            for frame_chunk_size in [None, 100, 10, 1]:
                for _AUDIO_REF_FP, _FEATS_REF_FP in zip(
                    [_SHORT_AUDIO_REF_FP, _LONG_AUDIO_REF_FP],
                    [_SHORT_FEATS_REF_FP, _LONG_FEATS_REF_FP],
                ):
                    audio = np.load(_AUDIO_REF_FP).reshape(1, -1)
                    audio = np.repeat(audio, batch_size, axis=0)
                    audio = torch.tensor(audio, requires_grad=False, device=self.device)
                    with torch.no_grad():
                        feats = (
                            self.extractor(audio, frame_chunk_size=frame_chunk_size)
                            .cpu()
                            .numpy()
                        )

                    feats_ref = np.load(_FEATS_REF_FP)[np.newaxis]
                    feats_ref = np.repeat(feats_ref, batch_size, axis=0)
                    # TODO: Fix shape here... How many frames do we expect 9585510 samples to have?
                    self.assertTrue(feats.shape[1] >= feats_ref.shape[1])

                    error = np.mean(np.abs(feats[:, : feats_ref.shape[1]] - feats_ref))
                    self.assertTrue(error < 0.012)


if __name__ == "__main__":
    unittest.main()

ddc_onset/spectral_test.py
chrisdonahue/ddc_onset
from typing import Optional

import numpy as np
import torch

from .spectral import SpectrogramExtractor
from .cnn import SpectrogramNormalizer, PlacementCNN
from .constants import SAMPLE_RATE, Difficulty

_MODULE_SINGLETONS = None


def compute_onset_salience(
    audio: np.ndarray,
    sr: int,
    device: Optional[torch.device] = None,
    difficulty=Difficulty.CHALLENGE,
):
    """Computes onset salience function from audio.

    Args:
        audio: Audio as float32 [num_samples].
        sr: Sample rate of audio. Resamples if not 44.1kHz.
        device: Device to run on. If None, run on CPU.
        difficulty: DDR difficulty label. For general onset detection, the default of Difficulty.CHALLENGE is recommended.
    Returns:
        Onset salience function as float32 [num_frames], frame rate of 100Hz.
    """
    # Check and maybe resample input
    if audio.dtype != np.float32:
        raise TypeError()
    if audio.ndim != 1:
        # TODO: Support multichannel audio and/or batched audio.
        raise ValueError()
    if sr != SAMPLE_RATE:
        try:
            import resampy
        except ImportError:
            raise Exception(
                "Resampy required to resample audio to 44.1kHz. Please install resampy with `pip install resampy`."
            )
        audio = resampy.resample(audio, sr, SAMPLE_RATE)

    # Load models
    global _MODULE_SINGLETONS
    if _MODULE_SINGLETONS is None:
        _MODULE_SINGLETONS = (
            SpectrogramExtractor(),
            SpectrogramNormalizer(),
            PlacementCNN(),
        )
        _MODULE_SINGLETONS = tuple(model.eval() for model in _MODULE_SINGLETONS)
    if device is not None:
        _MODULE_SINGLETONS = tuple(model.to(device) for model in _MODULE_SINGLETONS)
    extract, normalize, place = _MODULE_SINGLETONS

    # Predict
    with torch.no_grad():
        audio = torch.tensor(audio, device=device).view(1, -1)
        difficulties = torch.tensor(
            [difficulty.value], device=device, dtype=torch.int64
        )
        spectrogram = extract(audio)
        # TODO: Batch support
        onset_salience = place(normalize(spectrogram[0]), difficulties)[0]
        return onset_salience.cpu().numpy()


def find_peaks(onset_salience: np.ndarray):
    """Finds peaks in onset salience function."""
    try:
        from scipy.signal import argrelextrema
    except ImportError:
        raise Exception(
            "Scipy required for finding peaks. Please install scipy with `pip install scipy`."
        )
    onset_salience_smoothed = np.convolve(onset_salience, np.hamming(5), "same")
    peaks = argrelextrema(onset_salience_smoothed, np.greater_equal, order=1)[0]
    return peaks


def threshold_peaks(onset_salience: np.ndarray, peaks: np.ndarray, threshold: float):
    """Returns peaks above threshold."""
    return [i for i in peaks if onset_salience[i] >= threshold]

ddc_onset/util.py
chrisdonahue/ddc_onset
from setuptools import setup

setup(
    name="ddc_onset",
    version="0.1",
    description="Onset detector from Dance Dance Convolution in PyTorch",
    url="https://github.com/chrisdonahue/ddc_onset",
    author="Chris Donahue",
    packages=["ddc_onset"],
    package_data={"ddc_onset": ["weights/*"]},
    include_package_data=True,
    install_requires=[
        "numpy",
        "torch>=1.9.0",
    ],
)

setup.py
chrisdonahue/ddc_onset