hugo flores garcia
recovering from a gittastrophe
41b9d24
raw
history blame
7.52 kB
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)