Data processing

Datasets

class audiotools.data.datasets.AudioDataset(loaders: ~typing.Union[~audiotools.data.datasets.AudioLoader, ~typing.List[~audiotools.data.datasets.AudioLoader], ~typing.Dict[str, ~audiotools.data.datasets.AudioLoader]], sample_rate: int, n_examples: int = 1000, duration: float = 0.5, offset: ~typing.Optional[float] = None, loudness_cutoff: float = -40, num_channels: int = 1, transform: ~typing.Optional[~typing.Callable] = None, aligned: bool = False, shuffle_loaders: bool = False, matcher: ~typing.Callable = <function default_matcher>, without_replacement: bool = True)[source]

Bases: object

Loads audio from multiple loaders (with associated transforms) for a specified number of samples. Excerpts are drawn randomly of the specified duration, above a specified loudness threshold and are resampled on the fly to the desired sample rate (if it is different from the audio source sample rate).

This takes either a single AudioLoader object, a dictionary of AudioLoader objects, or a dictionary of AudioLoader objects. Each AudioLoader is called by the dataset, and the result is placed in the output dictionary. A transform can also be specified for the entire dataset, rather than for each specific loader. This transform can be applied to the output of all the loaders if desired.

AudioLoader objects can be specified as aligned, which means the loaders correspond to multitrack audio (e.g. a vocals, bass, drums, and other loader for multitrack music mixtures).

Parameters

loaders (Union[AudioLoader, List[AudioLoader], Dict[str, AudioLoader]]) – AudioLoaders to sample audio from.
sample_rate (int) – Desired sample rate.
n_examples (int, optional) – Number of examples (length of dataset), by default 1000
duration (float, optional) – Duration of audio samples, by default 0.5
loudness_cutoff (float, optional) – Loudness cutoff threshold for audio samples, by default -40
num_channels (int, optional) – Number of channels in output audio, by default 1
transform (Callable, optional) – Transform to instantiate alongside each dataset item, by default None
aligned (bool, optional) – Whether the loaders should be sampled in an aligned manner (e.g. same offset, duration, and matched file name), by default False
shuffle_loaders (bool, optional) – Whether to shuffle the loaders before sampling from them, by default False
matcher (Callable) – How to match files from adjacent audio lists (e.g. for a multitrack audio loader), by default uses the parent directory of each file.
without_replacement (bool) – Whether to choose files with or without replacement, by default True.

Examples

>>> from audiotools.data.datasets import AudioLoader
>>> from audiotools.data.datasets import AudioDataset
>>> from audiotools import transforms as tfm
>>> import numpy as np
>>>
>>> loaders = [
>>>     AudioLoader(
>>>         sources=[f"tests/audio/spk"],
>>>         transform=tfm.Equalizer(),
>>>         ext=["wav"],
>>>     )
>>>     for i in range(5)
>>> ]
>>>
>>> dataset = AudioDataset(
>>>     loaders = loaders,
>>>     sample_rate = 44100,
>>>     duration = 1.0,
>>>     transform = tfm.RescaleAudio(),
>>> )
>>>
>>> item = dataset[np.random.randint(len(dataset))]
>>>
>>> for i in range(len(loaders)):
>>>     item[i]["signal"] = loaders[i].transform(
>>>         item[i]["signal"], **item[i]["transform_args"]
>>>     )
>>>     item[i]["signal"].widget(i)
>>>
>>> mix = sum([item[i]["signal"] for i in range(len(loaders))])
>>> mix = dataset.transform(mix, **item["transform_args"])
>>> mix.widget("mix")

Below is an example of how one could load MUSDB multitrack data:

>>> import audiotools as at
>>> from pathlib import Path
>>> from audiotools import transforms as tfm
>>> import numpy as np
>>> import torch
>>>
>>> def build_dataset(
>>>     sample_rate: int = 44100,
>>>     duration: float = 5.0,
>>>     musdb_path: str = "~/.data/musdb/",
>>> ):
>>>     musdb_path = Path(musdb_path).expanduser()
>>>     loaders = {
>>>         src: at.datasets.AudioLoader(
>>>             sources=[musdb_path],
>>>             transform=tfm.Compose(
>>>                 tfm.VolumeNorm(("uniform", -20, -10)),
>>>                 tfm.Silence(prob=0.1),
>>>             ),
>>>             ext=[f"{src}.wav"],
>>>         )
>>>         for src in ["vocals", "bass", "drums", "other"]
>>>     }
>>>
>>>     dataset = at.datasets.AudioDataset(
>>>         loaders=loaders,
>>>         sample_rate=sample_rate,
>>>         duration=duration,
>>>         num_channels=1,
>>>         aligned=True,
>>>         transform=tfm.RescaleAudio(),
>>>         shuffle_loaders=True,
>>>     )
>>>     return dataset, list(loaders.keys())
>>>
>>> train_data, sources = build_dataset()
>>> dataloader = torch.utils.data.DataLoader(
>>>     train_data,
>>>     batch_size=16,
>>>     num_workers=0,
>>>     collate_fn=train_data.collate,
>>> )
>>> batch = next(iter(dataloader))
>>>
>>> for k in sources:
>>>     src = batch[k]
>>>     src["transformed"] = train_data.loaders[k].transform(
>>>         src["signal"].clone(), **src["transform_args"]
>>>     )
>>>
>>> mixture = sum(batch[k]["transformed"] for k in sources)
>>> mixture = train_data.transform(mixture, **batch["transform_args"])
>>>
>>> # Say a model takes the mix and gives back (n_batch, n_src, n_time).
>>> # Construct the targets:
>>> targets = at.AudioSignal.batch([batch[k]["transformed"] for k in sources], dim=1)

Similarly, here’s example code for loading Slakh data:

>>> import audiotools as at
>>> from pathlib import Path
>>> from audiotools import transforms as tfm
>>> import numpy as np
>>> import torch
>>> import glob
>>>
>>> def build_dataset(
>>>     sample_rate: int = 16000,
>>>     duration: float = 10.0,
>>>     slakh_path: str = "~/.data/slakh/",
>>> ):
>>>     slakh_path = Path(slakh_path).expanduser()
>>>
>>>     # Find the max number of sources in Slakh
>>>     src_names = [x.name for x in list(slakh_path.glob("**/*.wav"))  if "S" in str(x.name)]
>>>     n_sources = len(list(set(src_names)))
>>>
>>>     loaders = {
>>>         f"S{i:02d}": at.datasets.AudioLoader(
>>>             sources=[slakh_path],
>>>             transform=tfm.Compose(
>>>                 tfm.VolumeNorm(("uniform", -20, -10)),
>>>                 tfm.Silence(prob=0.1),
>>>             ),
>>>             ext=[f"S{i:02d}.wav"],
>>>         )
>>>         for i in range(n_sources)
>>>     }
>>>     dataset = at.datasets.AudioDataset(
>>>         loaders=loaders,
>>>         sample_rate=sample_rate,
>>>         duration=duration,
>>>         num_channels=1,
>>>         aligned=True,
>>>         transform=tfm.RescaleAudio(),
>>>         shuffle_loaders=False,
>>>     )
>>>
>>>     return dataset, list(loaders.keys())
>>>
>>> train_data, sources = build_dataset()
>>> dataloader = torch.utils.data.DataLoader(
>>>     train_data,
>>>     batch_size=16,
>>>     num_workers=0,
>>>     collate_fn=train_data.collate,
>>> )
>>> batch = next(iter(dataloader))
>>>
>>> for k in sources:
>>>     src = batch[k]
>>>     src["transformed"] = train_data.loaders[k].transform(
>>>         src["signal"].clone(), **src["transform_args"]
>>>     )
>>>
>>> mixture = sum(batch[k]["transformed"] for k in sources)
>>> mixture = train_data.transform(mixture, **batch["transform_args"])

static collate(list_of_dicts: Union[list, dict], n_splits: Optional[int] = None)[source]

Collates items drawn from this dataset. Uses audiotools.core.util.collate().

Parameters

list_of_dicts (Union[list, dict]) – Data drawn from each item.
n_splits (int) – Number of splits to make when creating the batches (split into sub-batches). Useful for things like gradient accumulation.

Returns

Dictionary of batched data.

Return type

dict

class audiotools.data.datasets.AudioLoader(sources: Optional[List[str]] = None, weights: Optional[List[float]] = None, transform: Optional[Callable] = None, relative_path: str = '', ext: List[str] = ['.wav', '.flac', '.mp3', '.mp4'], shuffle: bool = True, shuffle_state: int = 0)[source]

Bases: object

Loads audio endlessly from a list of audio sources containing paths to audio files. Audio sources can be folders full of audio files (which are found via file extension) or by providing a CSV file which contains paths to audio files.

Parameters

sources (List[str], optional) – Sources containing folders, or CSVs with paths to audio files, by default None
weights (List[float], optional) – Weights to sample audio files from each source, by default None
relative_path (str, optional) – Path audio should be loaded relative to, by default “”
transform (Callable, optional) – Transform to instantiate alongside audio sample, by default None
ext (List[str]) – List of extensions to find audio within each source by. Can also be a file name (e.g. “vocals.wav”). by default ['.wav', '.flac', '.mp3', '.mp4'].
shuffle (bool) – Whether to shuffle the files within the dataloader. Defaults to True.
shuffle_state (int) – State to use to seed the shuffle of the files.

class audiotools.data.datasets.ConcatDataset(datasets: list)[source]: Bases: AudioDataset

class audiotools.data.datasets.ResumableDistributedSampler(dataset, start_idx: Optional[int] = None, **kwargs)[source]

Bases: DistributedSampler

Distributed sampler that can be resumed from a given start index.

class audiotools.data.datasets.ResumableSequentialSampler(dataset, start_idx: Optional[int] = None, **kwargs)[source]

Bases: SequentialSampler

Sequential sampler that can be resumed from a given start index.

data_source: Sized

audiotools.data.datasets.align_lists(lists, matcher: ~typing.Callable = <function default_matcher>)[source]

audiotools.data.datasets.default_matcher(x, y)[source]

Preprocessing data

audiotools.data.preprocess.create_csv(audio_files: list, output_csv: Path, loudness: bool = False, data_path: Optional[str] = None)[source]

Converts a folder of audio files to a CSV file. If loudness = True, the output of this function will create a CSV file that looks something like:

path	loudness
daps/produced/f1_script1_produced.wav	-16.299999237060547
daps/produced/f1_script2_produced.wav	-16.600000381469727
daps/produced/f1_script3_produced.wav	-17.299999237060547
daps/produced/f1_script4_produced.wav	-16.100000381469727
daps/produced/f1_script5_produced.wav	-16.700000762939453
daps/produced/f3_script1_produced.wav	-16.5

Note

The paths above are written relative to the data_path argument which defaults to the environment variable PATH_TO_DATA if it isn’t passed to this function, and defaults to the empty string if that environment variable is not set.

You can produce a CSV file from a directory of audio files via:

>>> import audiotools
>>> directory = ...
>>> audio_files = audiotools.util.find_audio(directory)
>>> output_path = "train.csv"
>>> audiotools.data.preprocess.create_csv(
>>>     audio_files, output_csv, loudness=True
>>> )

Note that you can create empty rows in the CSV file by passing an empty string or None in the audio_files list. This is useful if you want to sync multiple CSV files in a multitrack setting. The loudness of these empty rows will be set to -inf.

Parameters

audio_files (list) – List of audio files.
output_csv (Path) – Output CSV, with each row containing the relative path of every file to data_path, if specified (defaults to None).
loudness (bool) – Compute loudness of entire file and store alongside path.

Transforms for data augmentation

class audiotools.data.transforms.BackgroundNoise(snr: tuple = ('uniform', 10.0, 30.0), sources: Optional[List[str]] = None, weights: Optional[List[float]] = None, eq_amount: tuple = ('const', 1.0), n_bands: int = 3, name: Optional[str] = None, prob: float = 1.0, loudness_cutoff: Optional[float] = None)[source]

path
room_tone/m6_script2_clean.wav
room_tone/m6_script2_cleanraw.wav
room_tone/m6_script2_ipad_balcony1.wav
room_tone/m6_script2_ipad_bedroom1.wav
room_tone/m6_script2_ipad_confroom1.wav
room_tone/m6_script2_ipad_confroom2.wav
room_tone/m6_script2_ipad_livingroom1.wav
room_tone/m6_script2_ipad_office1.wav