vampnet/beats.py

import json
import logging
import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from typing import List
from typing import Tuple
from typing import Union

import librosa
import torch
import numpy as np
from audiotools import AudioSignal


logging.basicConfig(level=logging.INFO)

###################
# beat sync utils #
###################

AGGREGATOR_REGISTRY = {
    "mean": np.mean,
    "median": np.median,
    "max": np.max,
    "min": np.min,
}


def list_aggregators() -> list:
    return list(AGGREGATOR_REGISTRY.keys())


@dataclass
class TimeSegment:
    start: float
    end: float

    @property
    def duration(self):
        return self.end - self.start

    def __str__(self) -> str:
        return f"{self.start} - {self.end}"

    def find_overlapping_segment(
        self, segments: List["TimeSegment"]
    ) -> Union["TimeSegment", None]:
        """Find the first segment that overlaps with this segment, or None if no segment overlaps"""
        for s in segments:
            if s.start <= self.start and s.end >= self.end:
                return s
        return None


def mkdir(path: Union[Path, str]) -> Path:
    p = Path(path)
    p.mkdir(parents=True, exist_ok=True)
    return p



###################
#    beat data    #
###################
@dataclass
class BeatSegment(TimeSegment):
    downbeat: bool = False  # if there's a downbeat on the start_time


class Beats:
    def __init__(self, beat_times, downbeat_times):
        if isinstance(beat_times, np.ndarray):
            beat_times = beat_times.tolist()
        if isinstance(downbeat_times, np.ndarray):
            downbeat_times = downbeat_times.tolist()
        self._beat_times = beat_times
        self._downbeat_times = downbeat_times
        self._use_downbeats = False

    def use_downbeats(self, use_downbeats: bool = True):
        """use downbeats instead of beats when calling beat_times"""
        self._use_downbeats = use_downbeats

    def beat_segments(self, signal: AudioSignal) -> List[BeatSegment]:
        """
        segments a song into time segments corresponding to beats.
        the first segment starts at 0 and ends at the first beat time.
        the last segment starts at the last beat time and ends at the end of the song.
        """
        beat_times = self._beat_times.copy()
        downbeat_times = self._downbeat_times
        beat_times.insert(0, 0)
        beat_times.append(signal.signal_duration)

        downbeat_ids = np.intersect1d(beat_times, downbeat_times, return_indices=True)[
            1
        ]
        is_downbeat = [
            True if i in downbeat_ids else False for i in range(len(beat_times))
        ]
        segments = [
            BeatSegment(start_time, end_time, downbeat)
            for start_time, end_time, downbeat in zip(
                beat_times[:-1], beat_times[1:], is_downbeat
            )
        ]
        return segments

    def get_beats(self) -> np.ndarray:
        """returns an array of beat times, in seconds
        if downbeats is True, returns an array of downbeat times, in seconds
        """
        return np.array(
            self._downbeat_times if self._use_downbeats else self._beat_times
        )

    @property
    def beat_times(self) -> np.ndarray:
        """return beat times"""
        return np.array(self._beat_times)

    @property
    def downbeat_times(self) -> np.ndarray:
        """return downbeat times"""
        return np.array(self._downbeat_times)

    def beat_times_to_feature_frames(
        self, signal: AudioSignal, features: np.ndarray
    ) -> np.ndarray:
        """convert beat times to frames, given an array of time-varying features"""
        beat_times = self.get_beats()
        beat_frames = (
            beat_times * signal.sample_rate / signal.signal_length * features.shape[-1]
        ).astype(np.int64)
        return beat_frames

    def sync_features(
        self, feature_frames: np.ndarray, features: np.ndarray, aggregate="median"
    ) -> np.ndarray:
        """sync features to beats"""
        if aggregate not in AGGREGATOR_REGISTRY:
            raise ValueError(f"unknown aggregation method {aggregate}")

        return librosa.util.sync(
            features, feature_frames, aggregate=AGGREGATOR_REGISTRY[aggregate]
        )

    def to_json(self) -> dict:
        """return beats and downbeats as json"""
        return {
            "beats": self._beat_times,
            "downbeats": self._downbeat_times,
            "use_downbeats": self._use_downbeats,
        }

    @classmethod
    def from_dict(cls, data: dict):
        """load beats and downbeats from json"""
        inst = cls(data["beats"], data["downbeats"])
        inst.use_downbeats(data["use_downbeats"])
        return inst

    def save(self, output_dir: Path):
        """save beats and downbeats to json"""
        mkdir(output_dir)
        with open(output_dir / "beats.json", "w") as f:
            json.dump(self.to_json(), f)

    @classmethod
    def load(cls, input_dir: Path):
        """load beats and downbeats from json"""
        beats_file = Path(input_dir) / "beats.json"
        with open(beats_file, "r") as f:
            data = json.load(f)
        return cls.from_dict(data)


###################
#  beat tracking  #
###################


class BeatTracker:
    def extract_beats(self, signal: AudioSignal) -> Tuple[np.ndarray, np.ndarray]:
        """extract beats from an audio signal"""
        raise NotImplementedError

    def __call__(self, signal: AudioSignal) -> Beats:
        """extract beats from an audio signal
        NOTE: if the first beat (and/or downbeat) is detected within the first 100ms of the audio,
        it is discarded. This is to avoid empty bins with no beat synced features in the first beat.
        Args:
            signal (AudioSignal): signal to beat track
        Returns:
            Tuple[np.ndarray, np.ndarray]: beats and downbeats
        """
        beats, downbeats = self.extract_beats(signal)
        return Beats(beats, downbeats)


class WaveBeat(BeatTracker):
    def __init__(self, ckpt_path: str = "checkpoints/wavebeat", device: str = "cpu"):
        from wavebeat.dstcn import dsTCNModel

        model = dsTCNModel.load_from_checkpoint(ckpt_path, map_location=torch.device(device))
        model.eval()

        self.device = device
        self.model = model

    def extract_beats(self, signal: AudioSignal) -> Tuple[np.ndarray, np.ndarray]:
        """returns beat and downbeat times, in  seconds"""
        # extract beats
        beats, downbeats = self.model.predict_beats_from_array(
            audio=signal.audio_data.squeeze(0),
            sr=signal.sample_rate,
            use_gpu=self.device != "cpu",
        )

        return beats, downbeats


class MadmomBeats(BeatTracker):
    def __init__(self):
        raise NotImplementedError

    def extract_beats(self, signal: AudioSignal) -> Tuple[np.ndarray, np.ndarray]:
        """returns beat and downbeat times, in  seconds"""
        pass


BEAT_TRACKER_REGISTRY = {
    "wavebeat": WaveBeat,
    "madmom": MadmomBeats,
}


def list_beat_trackers() -> list:
    return list(BEAT_TRACKER_REGISTRY.keys())


def load_beat_tracker(beat_tracker: str, **kwargs) -> BeatTracker:
    if beat_tracker not in BEAT_TRACKER_REGISTRY:
        raise ValueError(
            f"Unknown beat tracker {beat_tracker}. Available: {list_beat_trackers()}"
        )

    return BEAT_TRACKER_REGISTRY[beat_tracker](**kwargs)