/
/
/
1"""Log-mel spectrogram feature extractor for Beat This model."""
2
3from __future__ import annotations
4
5import asyncio
6import math
7from typing import TYPE_CHECKING, Any
8
9import numpy as np
10import torch
11import torchaudio
12
13if TYPE_CHECKING:
14 from collections.abc import Awaitable, Callable
15
16
17class AdvancedBeatFeatureExtractor:
18 """
19 Streaming log-mel extractor using torchaudio for Beat This compatibility.
20
21 Uses the same torchaudio.transforms.MelSpectrogram as beat_this.preprocessing.LogMelSpect:
22 - sample_rate=22050
23 - n_fft=1024
24 - hop_length=441
25 - f_min=30, f_max=11000
26 - n_mels=128
27 - mel_scale="slaney"
28 - normalized="frame_length"
29 - power=1
30 - center=True
31
32 Uses a sample-precise overlap approach that extracts exactly the frames
33 corresponding to each chunk's sample range. Delays the last few frames
34 of each chunk so they can be recomputed with real forward context from
35 the next chunk, avoiding reflect-padding artifacts at block boundaries.
36
37 Assumes input PCM is already at 22050 Hz mono.
38 """
39
40 def __init__(
41 self,
42 sample_rate: int = 22050,
43 n_fft: int = 1024,
44 hop_length: int = 441,
45 n_mels: int = 128,
46 fmin: float = 30.0,
47 fmax: float = 11000.0,
48 device: str = "cpu",
49 offload: Callable[..., Awaitable[Any]] | None = None,
50 ):
51 """
52 Initialize the feature extractor.
53
54 :param sample_rate: Audio sample rate (default 22050 Hz).
55 :param n_fft: FFT window size.
56 :param hop_length: Hop length between frames.
57 :param n_mels: Number of mel frequency bins.
58 :param fmin: Minimum frequency for mel filter.
59 :param fmax: Maximum frequency for mel filter.
60 :param device: Torch device to use.
61 :param offload: Awaitable runner for the blocking mel extraction. When given (the
62 provider passes its concurrency-bounded runner), it is used instead of a plain
63 asyncio.to_thread so the work counts against the host's analysis CPU cap.
64 """
65 self._offload = offload
66 self.n_fft = n_fft
67 self.hop_length = hop_length
68 self.sample_rate = sample_rate
69 self._device = device
70 self._n_mels = n_mels
71
72 # Number of frames to delay at the end of each chunk so they can be
73 # recomputed with real forward context from the next chunk.
74 # These frames need n_fft/2 samples of forward context that we don't
75 # have yet; delaying lets the next chunk provide it.
76 self._frames_to_delay = math.ceil((n_fft // 2) / hop_length)
77
78 # Track the total samples accumulated so far
79 self._total_samples = 0
80
81 # Buffer to hold samples from previous chunk needed for frame computation.
82 # Must be large enough to cover backward context for delayed frames:
83 # n_fft//2 (mel window half) + _frames_to_delay * hop_length (delayed range)
84 # + hop_length (hop-alignment margin).
85 self._keep_samples = n_fft + self._frames_to_delay * hop_length
86 self._prev_samples: np.ndarray | None = None
87
88 # Use torchaudio MelSpectrogram with center=True (standard beat_this approach)
89 self._mel_spec = torchaudio.transforms.MelSpectrogram(
90 sample_rate=sample_rate,
91 n_fft=n_fft,
92 hop_length=hop_length,
93 f_min=fmin,
94 f_max=fmax,
95 n_mels=n_mels,
96 mel_scale="slaney",
97 normalized="frame_length",
98 power=1,
99 center=True,
100 ).to(device)
101
102 # Track the last global frame index output by process_pcm
103 self._last_output_frame = -1
104
105 async def process_pcm(self, pcm: np.ndarray) -> np.ndarray:
106 """
107 Process a PCM chunk and return log-mel features.
108
109 :param pcm: Audio samples as float32 array.
110 :return: Log-mel features with shape (T, n_mels).
111 """
112
113 def _process_sync() -> np.ndarray:
114 chunk_start = self._total_samples
115 chunk_end = chunk_start + len(pcm)
116
117 # Determine which frames belong to this chunk.
118 # Frame j is centered at sample j * hop_length.
119 # If we have delayed frames from the previous chunk, start from there.
120 if chunk_start == 0:
121 first_frame = 0
122 elif self._last_output_frame >= 0:
123 first_frame = self._last_output_frame + 1
124 else:
125 first_frame = (chunk_start + self.hop_length - 1) // self.hop_length
126
127 # Last frame whose center is within this chunk
128 last_frame = (chunk_end - 1) // self.hop_length
129
130 # Delay the last frames so they can be recomputed with forward
131 # context from the next chunk. This avoids reflect-padding artifacts
132 # at block boundaries.
133 output_last_frame = last_frame - self._frames_to_delay
134
135 if output_last_frame < first_frame:
136 # Not enough frames to output anything yet; store samples and wait
137 if self._prev_samples is not None:
138 combined = np.concatenate([self._prev_samples, pcm])
139 else:
140 combined = pcm
141 self._prev_samples = combined[-self._keep_samples :].copy()
142 self._total_samples = chunk_end
143 return np.array([], dtype=np.float32).reshape(0, self._n_mels)
144
145 # Determine the audio range needed to compute these frames.
146 # Align audio_start to a hop_length boundary so segment frames
147 # correspond exactly to global frame positions.
148 needed_start = max(0, first_frame * self.hop_length - self.n_fft // 2)
149 audio_start = (needed_start // self.hop_length) * self.hop_length
150
151 # Build the audio segment to process
152 if self._prev_samples is not None and audio_start < chunk_start:
153 # We need some samples from the previous chunk
154 prev_needed = chunk_start - audio_start
155 prev_to_use = self._prev_samples[-prev_needed:]
156 audio_segment = np.concatenate([prev_to_use, pcm])
157 else:
158 audio_segment = pcm
159 audio_start = chunk_start
160
161 self._prev_samples = audio_segment[-self._keep_samples :].copy()
162
163 # Update total samples
164 self._total_samples = chunk_end
165
166 # Compute mel spectrogram
167 tensor = torch.from_numpy(audio_segment).to(self._device)
168 with torch.inference_mode():
169 mel = self._mel_spec(tensor)
170 log_mel = torch.log1p(1000.0 * mel)
171 features = log_mel.T.cpu().numpy()
172
173 # Extract frames for this chunk using integer arithmetic.
174 # Since audio_start is hop-aligned, segment frames map exactly to global frames.
175 segment_first_global_frame = audio_start // self.hop_length
176
177 start_in_segment = first_frame - segment_first_global_frame
178 end_in_segment = output_last_frame - segment_first_global_frame + 1
179
180 start_in_segment = max(0, start_in_segment)
181 end_in_segment = min(len(features), end_in_segment)
182
183 self._last_output_frame = output_last_frame
184
185 return features[start_in_segment:end_in_segment]
186
187 if self._offload is not None:
188 offloaded: np.ndarray = await self._offload(_process_sync)
189 return offloaded
190 return await asyncio.to_thread(_process_sync)
191
192 async def finalize(self) -> np.ndarray:
193 """
194 Flush delayed frames and process any remaining samples.
195
196 :return: Final log-mel features.
197 """
198
199 def _finalize_sync() -> np.ndarray:
200 if self._prev_samples is None or len(self._prev_samples) == 0:
201 return np.array([], dtype=np.float32).reshape(0, self._n_mels)
202
203 # mel_spec reflect-pad requires len > n_fft // 2
204 if len(self._prev_samples) <= self.n_fft // 2:
205 return np.array([], dtype=np.float32).reshape(0, self._n_mels)
206
207 # MelSpectrogram(center=True) produces 1 + total_samples // hop frames.
208 total_frames = 1 + self._total_samples // self.hop_length
209 extra_count = total_frames - self._last_output_frame - 1
210 if extra_count <= 0:
211 return np.array([], dtype=np.float32).reshape(0, self._n_mels)
212
213 tensor = torch.from_numpy(self._prev_samples).to(self._device)
214 with torch.inference_mode():
215 mel = self._mel_spec(tensor)
216 log_mel = torch.log1p(1000.0 * mel)
217 features = log_mel.T.cpu().numpy()
218
219 return features[-extra_count:]
220
221 if self._offload is not None:
222 offloaded: np.ndarray = await self._offload(_finalize_sync)
223 return offloaded
224 return await asyncio.to_thread(_finalize_sync)
225
226 def reset(self) -> None:
227 """Reset state for processing a new audio stream."""
228 self._total_samples = 0
229 self._prev_samples = None
230 self._last_output_frame = -1
231