wehub-resource-sync
167 lines
7.3 KiB
Python
167 lines
7.3 KiB
Python
import random
|
|
|
|
import numpy as np
|
|
import torch
|
|
from loguru import logger
|
|
from torch.utils.data import Dataset
|
|
from tqdm import tqdm
|
|
from yeaudio.audio import AudioSegment
|
|
from yeaudio.augmentation import ReverbPerturbAugmentor, SpecAugmentor
|
|
from yeaudio.augmentation import SpeedPerturbAugmentor, VolumePerturbAugmentor, NoisePerturbAugmentor
|
|
|
|
from mvector.data_utils.featurizer import AudioFeaturizer
|
|
|
|
|
|
class MVectorDataset(Dataset):
|
|
def __init__(self,
|
|
data_list_path,
|
|
audio_featurizer: AudioFeaturizer,
|
|
max_duration=3,
|
|
min_duration=0.5,
|
|
mode='train',
|
|
sample_rate=16000,
|
|
aug_conf=None,
|
|
num_speakers=None,
|
|
use_dB_normalization=True,
|
|
target_dB=-20):
|
|
"""音频数据加载器
|
|
|
|
Args:
|
|
data_list_path: 包含音频路径和标签的数据列表文件的路径
|
|
audio_featurizer: 声纹特征提取器
|
|
max_duration: 最长的音频长度，大于这个长度会裁剪掉
|
|
min_duration: 过滤最短的音频长度
|
|
aug_conf: 用于指定音频增强的配置
|
|
mode: 数据集模式。在训练模式下，数据集可能会进行一些数据增强的预处理
|
|
sample_rate: 采样率
|
|
num_speakers: 总说话人数量
|
|
use_dB_normalization: 是否对音频进行音量归一化
|
|
target_dB: 音量归一化的大小
|
|
"""
|
|
super(MVectorDataset, self).__init__()
|
|
assert mode in ['train', 'eval', 'extract_feature']
|
|
self.data_list_path = data_list_path
|
|
self.max_duration = max_duration
|
|
self.min_duration = min_duration
|
|
self.mode = mode
|
|
self._target_sample_rate = sample_rate
|
|
self._use_dB_normalization = use_dB_normalization
|
|
self._target_dB = target_dB
|
|
self.num_speakers = num_speakers
|
|
self.speed_augment = None
|
|
self.volume_augment = None
|
|
self.noise_augment = None
|
|
self.reverb_augment = None
|
|
self.spec_augment = None
|
|
# 获取特征器
|
|
self.audio_featurizer = audio_featurizer
|
|
# 获取特征裁剪的大小
|
|
self.max_feature_len = self.get_crop_feature_len()
|
|
# 获取数据列表
|
|
with open(self.data_list_path, 'r', encoding='utf-8') as f:
|
|
self.lines = f.readlines()
|
|
self.labels = [np.int64(line.strip().split('\t')[1]) for line in self.lines]
|
|
if mode == 'train' and aug_conf is not None:
|
|
# 获取数据增强器
|
|
self.get_augmentor(aug_conf)
|
|
# 评估模式下，数据列表需要排序
|
|
if self.mode == 'eval':
|
|
self.sort_list()
|
|
|
|
def __getitem__(self, idx):
|
|
# 分割数据文件路径和标签
|
|
data_path, spk_id = self.lines[idx].replace('\n', '').split('\t')
|
|
spk_id = int(spk_id)
|
|
# 如果后缀名为.npy的文件，那么直接读取
|
|
if data_path.endswith('.npy'):
|
|
feature = np.load(data_path)
|
|
if feature.shape[0] > self.max_feature_len:
|
|
crop_start = random.randint(0, feature.shape[0] - self.max_feature_len) if self.mode == 'train' else 0
|
|
feature = feature[crop_start:crop_start + self.max_feature_len, :]
|
|
feature = torch.tensor(feature, dtype=torch.float32)
|
|
else:
|
|
# 读取音频
|
|
audio_segment = AudioSegment.from_file(data_path)
|
|
# 数据太短不利于训练
|
|
if self.mode == 'train' or self.mode == 'extract_feature':
|
|
if audio_segment.duration < self.min_duration:
|
|
return self.__getitem__(idx + 1 if idx < len(self.lines) - 1 else 0)
|
|
# 重采样
|
|
if audio_segment.sample_rate != self._target_sample_rate:
|
|
audio_segment.resample(self._target_sample_rate)
|
|
# 音频增强
|
|
if self.mode == 'train':
|
|
audio_segment, spk_id = self.augment_audio(audio_segment, spk_id)
|
|
# decibel normalization
|
|
if self._use_dB_normalization:
|
|
audio_segment.normalize(target_db=self._target_dB)
|
|
# 裁剪需要的数据
|
|
if audio_segment.duration > self.max_duration:
|
|
audio_segment.crop(duration=self.max_duration, mode=self.mode)
|
|
samples = torch.tensor(audio_segment.samples, dtype=torch.float32)
|
|
try:
|
|
feature = self.audio_featurizer(samples)
|
|
except Exception as e:
|
|
logger.error(f"[{data_path}]特征提取失败，错误信息：{e}")
|
|
return self.__getitem__(idx + 1 if idx < len(self.lines) - 1 else 0)
|
|
feature = feature.squeeze(0)
|
|
if self.mode == 'train' and self.spec_augment is not None:
|
|
feature = self.spec_augment(feature.cpu().numpy())
|
|
feature = torch.tensor(feature, dtype=torch.float32)
|
|
spk_id = torch.tensor(spk_id, dtype=torch.int64)
|
|
return feature, spk_id
|
|
|
|
def __len__(self):
|
|
return len(self.lines)
|
|
|
|
# 获取特征裁剪的大小，对应max_duration音频提取特征后的长度
|
|
def get_crop_feature_len(self):
|
|
samples = torch.randn((1, int(self.max_duration * self._target_sample_rate)))
|
|
feature = self.audio_featurizer(samples).squeeze(0)
|
|
freq_len = feature.size(0)
|
|
return freq_len
|
|
|
|
# 数据列表需要排序
|
|
def sort_list(self):
|
|
lengths = []
|
|
for line in tqdm(self.lines, desc=f"对列表[{self.data_list_path}]进行长度排序"):
|
|
# 分割数据文件路径和标签
|
|
data_path, _ = line.split('\t')
|
|
if data_path.endswith('.npy'):
|
|
feature = np.load(data_path)
|
|
length = feature.shape[0]
|
|
lengths.append(length)
|
|
else:
|
|
# 读取音频
|
|
audio_segment = AudioSegment.from_file(data_path)
|
|
length = audio_segment.duration
|
|
lengths.append(length)
|
|
# 对长度排序并获取索引
|
|
sorted_indexes = np.argsort(lengths)
|
|
self.lines = [self.lines[i] for i in sorted_indexes]
|
|
|
|
# 获取数据增强器
|
|
def get_augmentor(self, aug_conf):
|
|
if aug_conf.speed is not None:
|
|
self.speed_augment = SpeedPerturbAugmentor(num_speakers=self.num_speakers, **aug_conf.speed)
|
|
if aug_conf.volume is not None:
|
|
self.volume_augment = VolumePerturbAugmentor(**aug_conf.volume)
|
|
if aug_conf.noise is not None:
|
|
self.noise_augment = NoisePerturbAugmentor(**aug_conf.noise)
|
|
if aug_conf.reverb is not None:
|
|
self.reverb_augment = ReverbPerturbAugmentor(**aug_conf.reverb)
|
|
if aug_conf.spec_aug is not None:
|
|
self.spec_augment = SpecAugmentor(**aug_conf.spec_aug)
|
|
|
|
# 音频增强
|
|
def augment_audio(self, audio_segment, spk_id):
|
|
if self.speed_augment is not None:
|
|
audio_segment, spk_id = self.speed_augment(audio_segment, spk_id)
|
|
if self.volume_augment is not None:
|
|
audio_segment = self.volume_augment(audio_segment)
|
|
if self.noise_augment is not None:
|
|
audio_segment = self.noise_augment(audio_segment)
|
|
if self.reverb_augment is not None:
|
|
audio_segment = self.reverb_augment(audio_segment)
|
|
return audio_segment, spk_id
|