File size: 13,381 Bytes
0d93e4e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
"""Dataset of audio with a simple description.
"""

from dataclasses import dataclass, fields, replace
import json
from pathlib import Path
import random
import typing as tp

import numpy as np
import torch

from .info_audio_dataset import (
    InfoAudioDataset,
    get_keyword_or_keyword_list
)
from ..modules.conditioners import (
    ConditioningAttributes,
    SegmentWithAttributes,
    WavCondition,
)


EPS = torch.finfo(torch.float32).eps
TARGET_LEVEL_LOWER = -35
TARGET_LEVEL_UPPER = -15


@dataclass
class SoundInfo(SegmentWithAttributes):
    """Segment info augmented with Sound metadata.
    """
    description: tp.Optional[str] = None
    self_wav: tp.Optional[torch.Tensor] = None

    @property
    def has_sound_meta(self) -> bool:
        return self.description is not None

    def to_condition_attributes(self) -> ConditioningAttributes:
        out = ConditioningAttributes()

        for _field in fields(self):
            key, value = _field.name, getattr(self, _field.name)
            if key == 'self_wav':
                out.wav[key] = value
            else:
                out.text[key] = value
        return out

    @staticmethod
    def attribute_getter(attribute):
        if attribute == 'description':
            preprocess_func = get_keyword_or_keyword_list
        else:
            preprocess_func = None
        return preprocess_func

    @classmethod
    def from_dict(cls, dictionary: dict, fields_required: bool = False):
        _dictionary: tp.Dict[str, tp.Any] = {}

        # allow a subset of attributes to not be loaded from the dictionary
        # these attributes may be populated later
        post_init_attributes = ['self_wav']

        for _field in fields(cls):
            if _field.name in post_init_attributes:
                continue
            elif _field.name not in dictionary:
                if fields_required:
                    raise KeyError(f"Unexpected missing key: {_field.name}")
            else:
                preprocess_func: tp.Optional[tp.Callable] = cls.attribute_getter(_field.name)
                value = dictionary[_field.name]
                if preprocess_func:
                    value = preprocess_func(value)
                _dictionary[_field.name] = value
        return cls(**_dictionary)


class SoundDataset(InfoAudioDataset):
    """Sound audio dataset: Audio dataset with environmental sound-specific metadata.

    Args:
        info_fields_required (bool): Whether all the mandatory metadata fields should be in the loaded metadata.
        external_metadata_source (tp.Optional[str]): Folder containing JSON metadata for the corresponding dataset.
            The metadata files contained in this folder are expected to match the stem of the audio file with
            a json extension.
        aug_p (float): Probability of performing audio mixing augmentation on the batch.
        mix_p (float): Proportion of batch items that are mixed together when applying audio mixing augmentation.
        mix_snr_low (int): Lowerbound for SNR value sampled for mixing augmentation.
        mix_snr_high (int): Upperbound for SNR value sampled for mixing augmentation.
        mix_min_overlap (float): Minimum overlap between audio files when performing mixing augmentation.
        kwargs: Additional arguments for AudioDataset.

    See `audiocraft.data.info_audio_dataset.InfoAudioDataset` for full initialization arguments.
    """
    def __init__(
        self,
        *args,
        info_fields_required: bool = True,
        external_metadata_source: tp.Optional[str] = None,
        aug_p: float = 0.,
        mix_p: float = 0.,
        mix_snr_low: int = -5,
        mix_snr_high: int = 5,
        mix_min_overlap: float = 0.5,
        **kwargs
    ):
        kwargs['return_info'] = True  # We require the info for each song of the dataset.
        super().__init__(*args, **kwargs)
        self.info_fields_required = info_fields_required
        self.external_metadata_source = external_metadata_source
        self.aug_p = aug_p
        self.mix_p = mix_p
        if self.aug_p > 0:
            assert self.mix_p > 0, "Expecting some mixing proportion mix_p if aug_p > 0"
            assert self.channels == 1, "SoundDataset with audio mixing considers only monophonic audio"
        self.mix_snr_low = mix_snr_low
        self.mix_snr_high = mix_snr_high
        self.mix_min_overlap = mix_min_overlap

    def _get_info_path(self, path: tp.Union[str, Path]) -> Path:
        """Get path of JSON with metadata (description, etc.).
        If there exists a JSON with the same name as 'path.name', then it will be used.
        Else, such JSON will be searched for in an external json source folder if it exists.
        """
        info_path = Path(path).with_suffix('.json')
        if Path(info_path).exists():
            return info_path
        elif self.external_metadata_source and (Path(self.external_metadata_source) / info_path.name).exists():
            return Path(self.external_metadata_source) / info_path.name
        else:
            raise Exception(f"Unable to find a metadata JSON for path: {path}")

    def __getitem__(self, index):
        wav, info = super().__getitem__(index)
        info_data = info.to_dict()
        info_path = self._get_info_path(info.meta.path)
        if Path(info_path).exists():
            with open(info_path, 'r') as json_file:
                sound_data = json.load(json_file)
                sound_data.update(info_data)
                sound_info = SoundInfo.from_dict(sound_data, fields_required=self.info_fields_required)
                # if there are multiple descriptions, sample one randomly
                if isinstance(sound_info.description, list):
                    sound_info.description = random.choice(sound_info.description)
        else:
            sound_info = SoundInfo.from_dict(info_data, fields_required=False)

        sound_info.self_wav = WavCondition(
            wav=wav[None], length=torch.tensor([info.n_frames]),
            sample_rate=[sound_info.sample_rate], path=[info.meta.path], seek_time=[info.seek_time])

        return wav, sound_info

    def collater(self, samples):
        # when training, audio mixing is performed in the collate function
        wav, sound_info = super().collater(samples)  # SoundDataset always returns infos
        if self.aug_p > 0:
            wav, sound_info = mix_samples(wav, sound_info, self.aug_p, self.mix_p,
                                          snr_low=self.mix_snr_low, snr_high=self.mix_snr_high,
                                          min_overlap=self.mix_min_overlap)
        return wav, sound_info


def rms_f(x: torch.Tensor) -> torch.Tensor:
    return (x ** 2).mean(1).pow(0.5)


def normalize(audio: torch.Tensor, target_level: int = -25) -> torch.Tensor:
    """Normalize the signal to the target level."""
    rms = rms_f(audio)
    scalar = 10 ** (target_level / 20) / (rms + EPS)
    audio = audio * scalar.unsqueeze(1)
    return audio


def is_clipped(audio: torch.Tensor, clipping_threshold: float = 0.99) -> torch.Tensor:
    return (abs(audio) > clipping_threshold).any(1)


def mix_pair(src: torch.Tensor, dst: torch.Tensor, min_overlap: float) -> torch.Tensor:
    start = random.randint(0, int(src.shape[1] * (1 - min_overlap)))
    remainder = src.shape[1] - start
    if dst.shape[1] > remainder:
        src[:, start:] = src[:, start:] + dst[:, :remainder]
    else:
        src[:, start:start+dst.shape[1]] = src[:, start:start+dst.shape[1]] + dst
    return src


def snr_mixer(clean: torch.Tensor, noise: torch.Tensor, snr: int, min_overlap: float,
              target_level: int = -25, clipping_threshold: float = 0.99) -> torch.Tensor:
    """Function to mix clean speech and noise at various SNR levels.

    Args:
        clean (torch.Tensor): Clean audio source to mix, of shape [B, T].
        noise (torch.Tensor): Noise audio source to mix, of shape [B, T].
        snr (int): SNR level when mixing.
        min_overlap (float): Minimum overlap between the two mixed sources.
        target_level (int): Gain level in dB.
        clipping_threshold (float): Threshold for clipping the audio.
    Returns:
        torch.Tensor: The mixed audio, of shape [B, T].
    """
    if clean.shape[1] > noise.shape[1]:
        noise = torch.nn.functional.pad(noise, (0, clean.shape[1] - noise.shape[1]))
    else:
        noise = noise[:, :clean.shape[1]]

    # normalizing to -25 dB FS
    clean = clean / (clean.max(1)[0].abs().unsqueeze(1) + EPS)
    clean = normalize(clean, target_level)
    rmsclean = rms_f(clean)

    noise = noise / (noise.max(1)[0].abs().unsqueeze(1) + EPS)
    noise = normalize(noise, target_level)
    rmsnoise = rms_f(noise)

    # set the noise level for a given SNR
    noisescalar = (rmsclean / (10 ** (snr / 20)) / (rmsnoise + EPS)).unsqueeze(1)
    noisenewlevel = noise * noisescalar

    # mix noise and clean speech
    noisyspeech = mix_pair(clean, noisenewlevel, min_overlap)

    # randomly select RMS value between -15 dBFS and -35 dBFS and normalize noisyspeech with that value
    # there is a chance of clipping that might happen with very less probability, which is not a major issue.
    noisy_rms_level = np.random.randint(TARGET_LEVEL_LOWER, TARGET_LEVEL_UPPER)
    rmsnoisy = rms_f(noisyspeech)
    scalarnoisy = (10 ** (noisy_rms_level / 20) / (rmsnoisy + EPS)).unsqueeze(1)
    noisyspeech = noisyspeech * scalarnoisy
    clean = clean * scalarnoisy
    noisenewlevel = noisenewlevel * scalarnoisy

    # final check to see if there are any amplitudes exceeding +/- 1. If so, normalize all the signals accordingly
    clipped = is_clipped(noisyspeech)
    if clipped.any():
        noisyspeech_maxamplevel = noisyspeech[clipped].max(1)[0].abs().unsqueeze(1) / (clipping_threshold - EPS)
        noisyspeech[clipped] = noisyspeech[clipped] / noisyspeech_maxamplevel

    return noisyspeech


def snr_mix(src: torch.Tensor, dst: torch.Tensor, snr_low: int, snr_high: int, min_overlap: float):
    if snr_low == snr_high:
        snr = snr_low
    else:
        snr = np.random.randint(snr_low, snr_high)
    mix = snr_mixer(src, dst, snr, min_overlap)
    return mix


def mix_text(src_text: str, dst_text: str):
    """Mix text from different sources by concatenating them."""
    if src_text == dst_text:
        return src_text
    return src_text + " " + dst_text


def mix_samples(wavs: torch.Tensor, infos: tp.List[SoundInfo], aug_p: float, mix_p: float,
                snr_low: int, snr_high: int, min_overlap: float):
    """Mix samples within a batch, summing the waveforms and concatenating the text infos.

    Args:
        wavs (torch.Tensor): Audio tensors of shape [B, C, T].
        infos (list[SoundInfo]): List of SoundInfo items corresponding to the audio.
        aug_p (float): Augmentation probability.
        mix_p (float): Proportion of items in the batch to mix (and merge) together.
        snr_low (int): Lowerbound for sampling SNR.
        snr_high (int): Upperbound for sampling SNR.
        min_overlap (float): Minimum overlap between mixed samples.
    Returns:
        tuple[torch.Tensor, list[SoundInfo]]: A tuple containing the mixed wavs
            and mixed SoundInfo for the given batch.
    """
    # no mixing to perform within the batch
    if mix_p == 0:
        return wavs, infos

    if random.uniform(0, 1) < aug_p:
        # perform all augmentations on waveforms as [B, T]
        # randomly picking pairs of audio to mix
        assert wavs.size(1) == 1, f"Mix samples requires monophonic audio but C={wavs.size(1)}"
        wavs = wavs.mean(dim=1, keepdim=False)
        B, T = wavs.shape
        k = int(mix_p * B)
        mixed_sources_idx = torch.randperm(B)[:k]
        mixed_targets_idx = torch.randperm(B)[:k]
        aug_wavs = snr_mix(
            wavs[mixed_sources_idx],
            wavs[mixed_targets_idx],
            snr_low,
            snr_high,
            min_overlap,
        )
        # mixing textual descriptions in metadata
        descriptions = [info.description for info in infos]
        aug_infos = []
        for i, j in zip(mixed_sources_idx, mixed_targets_idx):
            text = mix_text(descriptions[i], descriptions[j])
            m = replace(infos[i])
            m.description = text
            aug_infos.append(m)

        # back to [B, C, T]
        aug_wavs = aug_wavs.unsqueeze(1)
        assert aug_wavs.shape[0] > 0, "Samples mixing returned empty batch."
        assert aug_wavs.dim() == 3, f"Returned wav should be [B, C, T] but dim = {aug_wavs.dim()}"
        assert aug_wavs.shape[0] == len(aug_infos), "Mismatch between number of wavs and infos in the batch"

        return aug_wavs, aug_infos  # [B, C, T]
    else:
        # randomly pick samples in the batch to match
        # the batch size when performing audio mixing
        B, C, T = wavs.shape
        k = int(mix_p * B)
        wav_idx = torch.randperm(B)[:k]
        wavs = wavs[wav_idx]
        infos = [infos[i] for i in wav_idx]
        assert wavs.shape[0] == len(infos), "Mismatch between number of wavs and infos in the batch"

        return wavs, infos  # [B, C, T]