Skip to content

Commit 298b74e

Browse files
zhaotwzhaotw
committed
feat: add drum transcription and separation f
1 parent a350588 commit 298b74e

6 files changed

Lines changed: 671 additions & 2 deletions

File tree

‎README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ A collection of MIR open-source tools, dedicated to efficient inference tasks.
66
Designed to be lightweight with minimal package dependencies, ensuring easy installation and compatibility with various environments. This reduces the complexity of managing dependencies and helps avoid potential conflicts with other packages.
77

88
### Accelerated Multi-file Parallel Processing
9-
Significantly accelerating the analysis and processing of large datasets. By leveraging multi-threading and multi-processing techniques, and using CUDA for multi-GPU acceleration, MUSION can efficiently handle multiple files simultaneously, reducing the overall processing time.
9+
Significantly accelerating the analysis and processing of large datasets. By leveraging multi-threading and multi-processing techniques, and using ONNX TRT/CUDA for multi-GPU acceleration, MUSION can efficiently handle multiple files simultaneously, reducing the overall processing time.
1010

1111
### Unified and Concise UI
1212
Making it easy to use across different tools and applications. Whether you prefer using the Python interface, command-line interface (CLI), MUSION provides a consistent and intuitive experience for all users.
@@ -23,7 +23,7 @@ Separate songs into "drums", "bass", "other", "vocals"
2323
Detect chorus part for pop songs
2424

2525
[transcribe](musion/transcribe/README.md) automatic music transcription
26-
Transcribe piano or vocal audio to BEAT-ALIGNED MIDI file.
26+
Transcribe piano, drums or vocal audio to BEAT-ALIGNED MIDI file.
2727

2828
## Installation
2929
1. Using PyPI

‎musion/separate/drums.py‎

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
from typing import Optional
2+
import os
3+
4+
import torch
5+
import numpy as np
6+
7+
from musion.utils.ort_musion_base import OrtMusionBase
8+
from musion.utils.base import MusionPCM, FeatConfig
9+
10+
11+
MODULE_PATH = os.path.dirname(__file__)
12+
13+
class STFT:
14+
def __init__(self, device: torch.device):
15+
self.n_fft = 2048
16+
self.hop_length = 512
17+
self.device = device
18+
self.window = torch.hann_window(window_length=self.n_fft, periodic=True, device=device)
19+
self.dim_f = 1024
20+
21+
def __call__(self, x):
22+
batch_dims = x.shape[:-2]
23+
c, t = x.shape[-2:]
24+
x = x.reshape([-1, t])
25+
x = torch.stft(
26+
x,
27+
n_fft=self.n_fft,
28+
hop_length=self.hop_length,
29+
window=self.window,
30+
center=True,
31+
return_complex=True
32+
)
33+
x = torch.view_as_real(x)
34+
x = x.permute([0, 3, 1, 2])
35+
x = x.reshape([*batch_dims, c, 2, -1, x.shape[-1]]).reshape([*batch_dims, c * 2, -1, x.shape[-1]])
36+
return x[..., :self.dim_f, :]
37+
38+
def inverse(self, x):
39+
batch_dims = x.shape[:-3]
40+
c, f, t = x.shape[-3:]
41+
n = self.n_fft // 2 + 1
42+
f_pad = torch.zeros([*batch_dims, c, n - f, t], device=self.device, dtype=x.dtype)
43+
x = torch.cat([x, f_pad], -2)
44+
x = x.reshape([*batch_dims, c // 2, 2, n, t]).reshape([-1, 2, n, t])
45+
x = x.permute([0, 2, 3, 1])
46+
x = x[..., 0] + x[..., 1] * 1.j
47+
x = torch.istft(x, n_fft=self.n_fft, hop_length=self.hop_length, window=self.window, center=True)
48+
x = x.reshape([*batch_dims, 2, -1])
49+
return x
50+
51+
class _DrumsSeparate(OrtMusionBase):
52+
def __init__(self, device: str = None) -> None:
53+
OrtMusionBase.__init__(self,
54+
os.path.join(MODULE_PATH, 'separate_drums.onnx'),
55+
device,
56+
trt_fp16_enable=True)
57+
58+
self.stft = STFT(self.device)
59+
60+
@property
61+
def _feat_cfg(self) -> FeatConfig:
62+
return FeatConfig(
63+
mono=True,
64+
sample_rate=44100,
65+
)
66+
67+
def _getWindowingArray(self, window_size, fade_size):
68+
fadein = torch.linspace(0, 1, fade_size, device=self.device)
69+
fadeout = torch.linspace(1, 0, fade_size, device=self.device)
70+
window = torch.ones(window_size, device=self.device)
71+
window[-fade_size:] *= fadeout
72+
window[:fade_size] *= fadein
73+
return window
74+
75+
def demix_track(self, mix: np.ndarray):
76+
chunk_size = 130560
77+
N = 4
78+
fade_size = chunk_size // 10
79+
step = int(chunk_size // N)
80+
border = chunk_size - step
81+
batch_size = 1
82+
83+
full_length = mix.shape[-1]
84+
85+
# Do pad from the beginning and end to account floating window results better
86+
if full_length > 2 * border:
87+
mix = np.pad(mix, ((0, 0), (border, border)), mode='reflect')
88+
89+
# windowingArray crossfades at segment boundaries to mitigate clicking artifacts
90+
windowingArray = self._getWindowingArray(chunk_size, fade_size)
91+
92+
req_shape = (6,) + tuple(mix.shape)
93+
94+
result = torch.zeros(req_shape, dtype=torch.float32, device=self.device)
95+
counter = torch.zeros(req_shape, dtype=torch.float32, device=self.device)
96+
batch_data = []
97+
batch_locations = []
98+
99+
for i in range(0, mix.shape[1], step):
100+
chunk = mix[:, i:i+chunk_size]
101+
length = chunk.shape[-1]
102+
if length < chunk_size: # last chunk, pad
103+
pad_width = ((0,0), (0, chunk_size - length))
104+
if length > chunk_size // 2 + 1:
105+
chunk = np.pad(chunk, pad_width, mode='reflect')
106+
else:
107+
chunk = np.pad(chunk, pad_width, mode='constant')
108+
batch_data.append(torch.as_tensor(chunk, dtype=torch.float32, device=self.device))
109+
batch_locations.append((i, length))
110+
111+
if len(batch_data) >= batch_size:
112+
with torch.no_grad():
113+
arr = torch.stack(batch_data, axis=0)
114+
arr = self.stft(arr)
115+
x = self._predict_torch(arr)[0]
116+
if not torch.is_tensor(x):
117+
x = torch.as_tensor(x, device=self.device)
118+
x = self.stft.inverse(x)
119+
window = windowingArray
120+
if i - step == 0: # First audio chunk, no fadein
121+
window = window.clone()
122+
window[:fade_size] = 1
123+
elif i >= mix.shape[1]: # Last audio chunk, no fadeout
124+
window = window.clone()
125+
window[-fade_size:] = 1
126+
127+
for j in range(len(batch_locations)):
128+
start, l = batch_locations[j]
129+
result[..., start:start+l] += x[j][..., :l] * window[..., :l]
130+
counter[..., start:start+l] += window[..., :l]
131+
132+
batch_data = []
133+
batch_locations = []
134+
135+
estimated_sources = result / counter
136+
estimated_sources = torch.nan_to_num(estimated_sources, nan=0.0)
137+
138+
if full_length > 2 * border:
139+
# Remove pad
140+
estimated_sources = estimated_sources[..., border:-border]
141+
142+
estimated_sources = estimated_sources.detach().cpu().numpy()
143+
return {k: v for k, v in zip(['kick', 'snare', 'toms', 'hh', 'ride', 'crash'], estimated_sources)}
144+
145+
def _process(self, audio_path: Optional[str] = None, pcm: Optional[MusionPCM] = None) -> dict:
146+
mix = self._load_pcm(audio_path, pcm).samples[0]
147+
mix *= 10 ** (9 / 20) # 9dB gain
148+
149+
mix = np.stack([mix, mix], axis=-1) # Convert mono to stereo [time, channels]
150+
151+
res = self.demix_track(mix.T)
152+
153+
return res
154+
155+
@property
156+
def result_keys(self):
157+
return ['wav']

‎musion/transcribe/README.md‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,3 +21,15 @@ Return:
2121
{
2222
'vocals': [ [note start time, note end time, note pitch] x n ]
2323
}
24+
25+
3. Drums Transcription
26+
Encapsulation for [DrumTranscription](https://github.com/xavriley/DrumTranscription)
27+
28+
Input:
29+
audio file path or audio pcm
30+
31+
Return:
32+
{
33+
'mid': mido.MidiFile
34+
}
35+

‎musion/transcribe/drums/drums.py‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
import os
2+
from typing import Optional, Dict, List
3+
4+
import numpy as np
5+
import pretty_midi
6+
import madmom
7+
from madmom.audio.filters import LogarithmicFilterbank
8+
from madmom.audio.signal import FramedSignalProcessor, SignalProcessor
9+
from madmom.audio.spectrogram import LogarithmicFilteredSpectrogramProcessor
10+
from madmom.audio.stft import ShortTimeFourierTransformProcessor
11+
from madmom.processors import SequentialProcessor
12+
13+
from musion.utils.base import FeatConfig, MusionPCM
14+
from musion.utils.ort_musion_base import OrtMusionBase
15+
from musion.transcribe.base import *
16+
from musion.transcribe.drums.velocity import estimate_velocity
17+
from musion.separate.drums import _DrumsSeparate
18+
from musion.transcribe.midi_utils import MIDI_RESOLUTION
19+
20+
MODULE_PATH = os.path.dirname(__file__)
21+
22+
# Drum MIDI note numbers: ["BD", "SD", "TT", "HH", "CY+RD"]
23+
LABELS_5: List[int] = [35, 38, 47, 42, 49]
24+
FPS: int = 100
25+
DEFAULT_NOTE_DURATION: float = 0.01 # seconds
26+
DEFAULT_VELOCITY: int = 100 # Will be refined by estimate_velocity
27+
28+
class _DrumsTranscribe(TranscribeBase, OrtMusionBase):
29+
def __init__(self, device: str = None) -> None:
30+
TranscribeBase.__init__(self, device)
31+
OrtMusionBase.__init__(self,
32+
os.path.join(MODULE_PATH, 'transcribe_drums.onnx'),
33+
device)
34+
35+
frameSize = self._feat_cfg.n_fft
36+
audio_sample_rate = self._feat_cfg.sample_rate
37+
38+
sig = SignalProcessor(num_channels=1, sample_rate=audio_sample_rate)
39+
frames = FramedSignalProcessor(frame_size=frameSize, fps=FPS)
40+
stft = ShortTimeFourierTransformProcessor()
41+
spec = LogarithmicFilteredSpectrogramProcessor(
42+
num_channels=1,
43+
sample_rate=audio_sample_rate,
44+
filterbank=LogarithmicFilterbank,
45+
frame_size=frameSize,
46+
fps=FPS,
47+
num_bands=12,
48+
fmin=20,
49+
fmax=20000,
50+
norm_filters=True,
51+
)
52+
53+
self.pre_processor = SequentialProcessor((sig, frames, stft, spec))
54+
55+
# Peak picking thresholds for each drum type
56+
peak_thresholds: List[float] = [0.22, 0.24, 0.32, 0.22, 0.2]
57+
self.processors = [
58+
madmom.features.notes.NoteOnsetPeakPickingProcessor(
59+
threshold=t, smooth=0, pre_avg=0.1, post_avg=0.01, pre_max=0.02, post_max=0.01, combine=0.02, fps=FPS)
60+
for t in peak_thresholds
61+
]
62+
63+
self.separate_drums = _DrumsSeparate(self.device)
64+
65+
@property
66+
def _feat_cfg(self) -> FeatConfig:
67+
return FeatConfig(
68+
mono=True,
69+
sample_rate=44100,
70+
n_fft=2048,
71+
hop_length=441,
72+
f_min=20,
73+
f_max=20000,
74+
)
75+
76+
def _process(self, audio_path: Optional[str] = None, pcm: Optional[MusionPCM] = None) -> Dict[str, mido.MidiFile]:
77+
audio = self.pre_processor(audio_path)
78+
# Reshape to add channel dimension: [num_frames, num_bands] -> [num_frames, num_bands, 1]
79+
audio = audio.reshape((audio.shape[0], audio.shape[1], 1))
80+
81+
raw_probs = self.predict(audio)[0] # [num_frames, num_labels]
82+
83+
midi = pretty_midi.PrettyMIDI(resolution=MIDI_RESOLUTION)
84+
instrument = pretty_midi.Instrument(program=0, is_drum=True, name="Drums")
85+
midi.instruments.append(instrument)
86+
87+
for i, processor in enumerate(self.processors):
88+
# Reshape probability array for peak picking: [num_frames] -> [num_frames, 1]
89+
prob_array = raw_probs[:, i].reshape(-1, 1)
90+
peaks = processor.process(prob_array) # Shape: (num_peaks, 2) with (time, pitch)
91+
92+
for onset_time in peaks[:, 0]:
93+
note = pretty_midi.Note(
94+
velocity=DEFAULT_VELOCITY,
95+
pitch=LABELS_5[i],
96+
start=onset_time,
97+
end=onset_time + DEFAULT_NOTE_DURATION
98+
)
99+
instrument.notes.append(note)
100+
101+
drum_parts_separation_res = self.separate_drums(audio_path=audio_path)
102+
estimate_velocity(midi, drum_parts_separation_res)
103+
midi = self._align_midi_with_beats(midi, audio_path)
104+
105+
return {'mid': midi}
106+
107+
def predict(self, audio: np.ndarray, limit_input_size: int = 60000) -> np.ndarray:
108+
"""
109+
Predict drum transcription probabilities for audio input.
110+
111+
For RNN models, uses overlapping windows with warmup sequences to handle long audio.
112+
113+
Args:
114+
audio: Input audio features with shape [num_frames, num_bands, num_channels]
115+
limit_input_size: Maximum window size before potential segmentation fault.
116+
If input is larger, it's split into overlapping windows.
117+
118+
Returns:
119+
Probability array with shape [num_frames, num_labels]
120+
121+
TODO: Identify the error "Computed output size would be negative: -2
122+
[input_size: 0, effective_filter_size: 3, stride: 1]"
123+
"""
124+
window_size = limit_input_size
125+
warmup_size = 412
126+
step_size = window_size - 2 * warmup_size
127+
128+
predictions: List[np.ndarray] = []
129+
# Split into overlapping windows if input is too large
130+
window_indices = range(0, len(audio) - warmup_size, step_size)
131+
132+
for sample_idx in window_indices:
133+
window_audio = audio[sample_idx : sample_idx + window_size]
134+
# Add batch dimension: [num_frames, num_bands, num_channels] -> [1, num_frames, num_bands, num_channels]
135+
window_audio_batch = window_audio.reshape((1,) + window_audio.shape)
136+
prediction = self._predict([window_audio_batch])[0]
137+
138+
if len(window_indices) == 1:
139+
# Single window: no warmup needed
140+
predictions.append(prediction)
141+
elif sample_idx == 0:
142+
# First window: no warmup at beginning, only at end
143+
predictions.append(prediction[:window_size - warmup_size])
144+
elif sample_idx == window_indices[-1]:
145+
# Last window: warmup at beginning, none at end
146+
predictions.append(prediction[warmup_size:])
147+
else:
148+
# Middle windows: warmup at both beginning and end
149+
predictions.append(prediction[warmup_size:][:step_size])
150+
151+
return np.concatenate(predictions)

0 commit comments

Comments
 (0)