File size: 2,036 Bytes
162fa19 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | import torch
import numpy as np
from config import SharedConfig
from typing import Union, Optional
class Segmenter:
def __init__(self, config: SharedConfig):
self.config = config
def chunk_notes(self, notes: np.ndarray):
# create a mock audio
audio_length = notes[:, 1].max() # get the max offset
n_chunks = int(audio_length // self.config.input_sec) + 1
notes_chunks = [[] for _ in range(n_chunks)]
for note in notes:
onset, offset, pitch, velocity = note
on_idx = int(onset // self.config.input_sec)
off_idx = int(offset // self.config.input_sec)
onset = onset % self.config.input_sec
offset = offset % self.config.input_sec
if on_idx == off_idx:
notes_chunks[on_idx].append(np.array([onset, offset, pitch, velocity]))
return notes_chunks
def chunk_audio_and_notes(self, audio: torch.FloatTensor, notes: np.ndarray, audio_file: Optional[str] = None):
assert audio.ndim == 1, "audio must be a 1D tensor"
audio_chunks = list(audio.split(int(self.config.input_sec * self.config.sample_rate), dim=0))
notes_chunks = [[] for _ in range(len(audio_chunks))]
for note in notes:
onset, offset, pitch, velocity = note
original_onset, original_offset = onset, offset
on_idx = int(onset // self.config.input_sec)
off_idx = int(offset // self.config.input_sec)
onset = onset % self.config.input_sec
offset = offset % self.config.input_sec
if on_idx > len(notes_chunks) - 1:
raise ValueError(
f"onset index {on_idx} is greater than the number of notes chunks {len(notes_chunks)}. This means that annotations are not consistent with the audio length."
)
if on_idx == off_idx:
notes_chunks[on_idx].append(np.array([onset, offset, pitch, velocity]))
return audio_chunks, notes_chunks
|