Spaces:
Running on Zero
Running on Zero
File size: 4,969 Bytes
ff7b988 | 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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | import librosa
import math
import numpy as np
import jukebox.utils.dist_adapter as dist
from torch.utils.data import Dataset
from jukebox.utils.dist_utils import print_all
from jukebox.utils.io import get_duration_sec, load_audio
from jukebox.data.labels import Labeller
class FilesAudioDataset(Dataset):
def __init__(self, hps):
super().__init__()
self.sr = hps.sr
self.channels = hps.channels
self.min_duration = hps.min_duration or math.ceil(hps.sample_length / hps.sr)
self.max_duration = hps.max_duration or math.inf
self.sample_length = hps.sample_length
assert hps.sample_length / hps.sr < self.min_duration, f'Sample length {hps.sample_length} per sr {hps.sr} ({hps.sample_length / hps.sr:.2f}) should be shorter than min duration {self.min_duration}'
self.aug_shift = hps.aug_shift
self.labels = hps.labels
self.init_dataset(hps)
def filter(self, files, durations):
# Remove files too short or too long
keep = []
for i in range(len(files)):
if durations[i] / self.sr < self.min_duration:
continue
if durations[i] / self.sr >= self.max_duration:
continue
keep.append(i)
print_all(f'self.sr={self.sr}, min: {self.min_duration}, max: {self.max_duration}')
print_all(f"Keeping {len(keep)} of {len(files)} files")
self.files = [files[i] for i in keep]
self.durations = [int(durations[i]) for i in keep]
self.cumsum = np.cumsum(self.durations)
def init_dataset(self, hps):
# Load list of files and starts/durations
files = librosa.util.find_files(f'{hps.audio_files_dir}', ['mp3', 'opus', 'm4a', 'aac', 'wav'])
print_all(f"Found {len(files)} files. Getting durations")
cache = dist.get_rank() % 8 == 0 if dist.is_available() else True
durations = np.array([get_duration_sec(file, cache=cache) * self.sr for file in files]) # Could be approximate
self.filter(files, durations)
if self.labels:
self.labeller = Labeller(hps.max_bow_genre_size, hps.n_tokens, self.sample_length, v3=hps.labels_v3)
def get_index_offset(self, item):
# For a given dataset item and shift, return song index and offset within song
half_interval = self.sample_length//2
shift = np.random.randint(-half_interval, half_interval) if self.aug_shift else 0
offset = item * self.sample_length + shift # Note we centred shifts, so adding now
midpoint = offset + half_interval
assert 0 <= midpoint < self.cumsum[-1], f'Midpoint {midpoint} of item beyond total length {self.cumsum[-1]}'
index = np.searchsorted(self.cumsum, midpoint) # index <-> midpoint of interval lies in this song
start, end = self.cumsum[index - 1] if index > 0 else 0.0, self.cumsum[index] # start and end of current song
assert start <= midpoint <= end, f"Midpoint {midpoint} not inside interval [{start}, {end}] for index {index}"
if offset > end - self.sample_length: # Going over song
offset = max(start, offset - half_interval) # Now should fit
elif offset < start: # Going under song
offset = min(end - self.sample_length, offset + half_interval) # Now should fit
assert start <= offset <= end - self.sample_length, f"Offset {offset} not in [{start}, {end - self.sample_length}]. End: {end}, SL: {self.sample_length}, Index: {index}"
offset = offset - start
return index, offset
def get_metadata(self, filename, test):
"""
Insert metadata loading code for your dataset here.
If artist/genre labels are different from provided artist/genre lists,
update labeller accordingly.
Returns:
(artist, genre, full_lyrics) of type (str, str, str). For
example, ("unknown", "classical", "") could be a metadata for a
piano piece.
"""
return None, None, None
def get_song_chunk(self, index, offset, test=False):
filename, total_length = self.files[index], self.durations[index]
data, sr = load_audio(filename, sr=self.sr, offset=offset, duration=self.sample_length)
assert data.shape == (self.channels, self.sample_length), f'Expected {(self.channels, self.sample_length)}, got {data.shape}'
if self.labels:
artist, genre, lyrics = self.get_metadata(filename, test)
labels = self.labeller.get_label(artist, genre, lyrics, total_length, offset)
return data.T, labels['y']
else:
return data.T
def get_item(self, item, test=False):
index, offset = self.get_index_offset(item)
return self.get_song_chunk(index, offset, test)
def __len__(self):
return int(np.floor(self.cumsum[-1] / self.sample_length))
def __getitem__(self, item):
return self.get_item(item)
|