58 lines
2.5 KiB
Python
58 lines
2.5 KiB
Python
"""Contains the noise perturb augmentation model."""
|
||
import os
|
||
import random
|
||
|
||
import numpy as np
|
||
|
||
from mvector.data_utils.augmentor.base import AugmentorBase
|
||
from mvector.data_utils.audio import AudioSegment
|
||
|
||
|
||
class NoisePerturbAugmentor(AugmentorBase):
|
||
"""用于添加背景噪声的增强模型
|
||
|
||
:param min_snr_dB: Minimal signal noise ratio, in decibels.
|
||
:type min_snr_dB: float
|
||
:param max_snr_dB: Maximal signal noise ratio, in decibels.
|
||
:type max_snr_dB: float
|
||
:param repetition: repetition noise sum
|
||
:type repetition: int
|
||
:param noise_dir: noise audio file dir.
|
||
:type noise_dir: str
|
||
"""
|
||
|
||
def __init__(self, min_snr_dB, max_snr_dB, repetition, noise_dir):
|
||
self._min_snr_dB = min_snr_dB
|
||
self._max_snr_dB = max_snr_dB
|
||
self.repetition = repetition
|
||
self.noises_path = []
|
||
if os.path.exists(noise_dir):
|
||
for file in os.listdir(noise_dir):
|
||
self.noises_path.append(os.path.join(noise_dir, file))
|
||
|
||
def transform_audio(self, audio_segment: AudioSegment):
|
||
"""Add background noise audio.
|
||
|
||
Note that this is an in-place transformation.
|
||
|
||
:param audio_segment: Audio segment to add effects to.
|
||
:type audio_segment: AudioSegmenet
|
||
"""
|
||
if len(self.noises_path) > 0:
|
||
for _ in range(random.randint(1, self.repetition)):
|
||
# 随机选择一个noises_path中的一个
|
||
noise_path = random.sample(self.noises_path, 1)[0]
|
||
# 读取噪声音频
|
||
noise_segment = AudioSegment.from_file(noise_path)
|
||
# 如果噪声采样率不等于audio_segment的采样率，则重采样
|
||
if noise_segment.sample_rate != audio_segment.sample_rate:
|
||
noise_segment.resample(audio_segment.sample_rate)
|
||
# 随机生成snr_dB的值
|
||
snr_dB = random.uniform(self._min_snr_dB, self._max_snr_dB)
|
||
# 如果噪声的长度小于audio_segment的长度，则将噪声的前面的部分填充噪声末尾补长
|
||
if noise_segment.duration < audio_segment.duration:
|
||
diff_duration = audio_segment.num_samples - noise_segment.num_samples
|
||
noise_segment._samples = np.pad(noise_segment.samples, (0, diff_duration), 'wrap')
|
||
# 将噪声添加到audio_segment中，并将snr_dB调整到最小值和最大值之间
|
||
audio_segment.add_noise(noise_segment, snr_dB)
|