Spaces:
Runtime error
Runtime error
File size: 6,719 Bytes
efcaf28 |
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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 |
import os
import h5py
import numpy as np
from sortedcontainers import SortedList
from torch.utils.data import Dataset
from tqdm import tqdm
from data.utils import load
class SeparationDataset(Dataset):
def __init__(self, dataset, partition, instruments, sr, channels, shapes, random_hops, hdf_dir, audio_transform=None, in_memory=False):
'''
Initialises a source separation dataset
:param data: HDF audio data object
:param input_size: Number of input samples for each example
:param context_front: Number of extra context samples to prepend to input
:param context_back: NUmber of extra context samples to append to input
:param hop_size: Skip hop_size - 1 sample positions in the audio for each example (subsampling the audio)
:param random_hops: If False, sample examples evenly from whole audio signal according to hop_size parameter. If True, randomly sample a position from the audio
'''
super(SeparationDataset, self).__init__()
self.hdf_dataset = None
os.makedirs(hdf_dir, exist_ok=True)
self.hdf_dir = os.path.join(hdf_dir, partition + ".hdf5")
self.random_hops = random_hops
self.sr = sr
self.channels = channels
self.shapes = shapes
self.audio_transform = audio_transform
self.in_memory = in_memory
self.instruments = instruments
# PREPARE HDF FILE
# Check if HDF file exists already
if not os.path.exists(self.hdf_dir):
# Create folder if it did not exist before
if not os.path.exists(hdf_dir):
os.makedirs(hdf_dir)
# Create HDF file
with h5py.File(self.hdf_dir, "w") as f:
f.attrs["sr"] = sr
f.attrs["channels"] = channels
f.attrs["instruments"] = instruments
print("Adding audio files to dataset (preprocessing)...")
for idx, example in enumerate(tqdm(dataset[partition])):
# Load mix
mix_audio, _ = load(example["mix"], sr=self.sr, mono=(self.channels == 1))
source_audios = []
for source in instruments:
# In this case, read in audio and convert to target sampling rate
source_audio, _ = load(example[source], sr=self.sr, mono=(self.channels == 1))
source_audios.append(source_audio)
source_audios = np.concatenate(source_audios, axis=0)
assert(source_audios.shape[1] == mix_audio.shape[1])
# Add to HDF5 file
grp = f.create_group(str(idx))
grp.create_dataset("inputs", shape=mix_audio.shape, dtype=mix_audio.dtype, data=mix_audio)
grp.create_dataset("targets", shape=source_audios.shape, dtype=source_audios.dtype, data=source_audios)
grp.attrs["length"] = mix_audio.shape[1]
grp.attrs["target_length"] = source_audios.shape[1]
# In that case, check whether sr and channels are complying with the audio in the HDF file, otherwise raise error
with h5py.File(self.hdf_dir, "r") as f:
if f.attrs["sr"] != sr or \
f.attrs["channels"] != channels or \
list(f.attrs["instruments"]) != instruments:
raise ValueError(
"Tried to load existing HDF file, but sampling rate and channel or instruments are not as expected. Did you load an out-dated HDF file?")
# HDF FILE READY
# SET SAMPLING POSITIONS
# Go through HDF and collect lengths of all audio files
with h5py.File(self.hdf_dir, "r") as f:
lengths = [f[str(song_idx)].attrs["target_length"] for song_idx in range(len(f))]
# Subtract input_size from lengths and divide by hop size to determine number of starting positions
lengths = [(l // self.shapes["output_frames"]) + 1 for l in lengths]
self.start_pos = SortedList(np.cumsum(lengths))
self.length = self.start_pos[-1]
def __getitem__(self, index):
# Open HDF5
if self.hdf_dataset is None:
driver = "core" if self.in_memory else None # Load HDF5 fully into memory if desired
self.hdf_dataset = h5py.File(self.hdf_dir, 'r', driver=driver)
# Find out which slice of targets we want to read
audio_idx = self.start_pos.bisect_right(index)
if audio_idx > 0:
index = index - self.start_pos[audio_idx - 1]
# Check length of audio signal
audio_length = self.hdf_dataset[str(audio_idx)].attrs["length"]
target_length = self.hdf_dataset[str(audio_idx)].attrs["target_length"]
# Determine position where to start targets
if self.random_hops:
start_target_pos = np.random.randint(0, max(target_length - self.shapes["output_frames"] + 1, 1))
else:
# Map item index to sample position within song
start_target_pos = index * self.shapes["output_frames"]
# READ INPUTS
# Check front padding
start_pos = start_target_pos - self.shapes["output_start_frame"]
if start_pos < 0:
# Pad manually since audio signal was too short
pad_front = abs(start_pos)
start_pos = 0
else:
pad_front = 0
# Check back padding
end_pos = start_target_pos - self.shapes["output_start_frame"] + self.shapes["input_frames"]
if end_pos > audio_length:
# Pad manually since audio signal was too short
pad_back = end_pos - audio_length
end_pos = audio_length
else:
pad_back = 0
# Read and return
audio = self.hdf_dataset[str(audio_idx)]["inputs"][:, start_pos:end_pos].astype(np.float32)
if pad_front > 0 or pad_back > 0:
audio = np.pad(audio, [(0, 0), (pad_front, pad_back)], mode="constant", constant_values=0.0)
targets = self.hdf_dataset[str(audio_idx)]["targets"][:, start_pos:end_pos].astype(np.float32)
if pad_front > 0 or pad_back > 0:
targets = np.pad(targets, [(0, 0), (pad_front, pad_back)], mode="constant", constant_values=0.0)
targets = {inst : targets[idx*self.channels:(idx+1)*self.channels] for idx, inst in enumerate(self.instruments)}
if hasattr(self, "audio_transform") and self.audio_transform is not None:
audio, targets = self.audio_transform(audio, targets)
return audio, targets
def __len__(self):
return self.length |