74 lines
3.0 KiB
Python
74 lines
3.0 KiB
Python
import numpy as np
|
||
from torch.utils.data import Dataset
|
||
|
||
from mvector.data_utils.audio import AudioSegment
|
||
from mvector.data_utils.augmentor.augmentation import AugmentationPipeline
|
||
from mvector.utils.logger import setup_logger
|
||
|
||
logger = setup_logger(__name__)
|
||
|
||
|
||
class CustomDataset(Dataset):
|
||
def __init__(self,
|
||
data_list_path,
|
||
do_vad=True,
|
||
max_duration=6,
|
||
min_duration=0.5,
|
||
augmentation_config='{}',
|
||
mode='train',
|
||
sample_rate=16000,
|
||
use_dB_normalization=True,
|
||
target_dB=-20):
|
||
"""音频数据加载器
|
||
|
||
Args:
|
||
data_list_path: 包含音频路径和标签的数据列表文件的路径
|
||
do_vad: 是否对音频进行语音活动检测（VAD）来裁剪静音部分
|
||
max_duration: 最长的音频长度，大于这个长度会裁剪掉
|
||
min_duration: 过滤最短的音频长度
|
||
augmentation_config: 用于指定音频增强的配置
|
||
mode: 数据集模式。在训练模式下，数据集可能会进行一些数据增强的预处理
|
||
sample_rate: 采样率
|
||
use_dB_normalization: 是否对音频进行音量归一化
|
||
target_dB: 音量归一化的大小
|
||
"""
|
||
super(CustomDataset, self).__init__()
|
||
self.do_vad = do_vad
|
||
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._augmentation_pipeline = AugmentationPipeline(augmentation_config=augmentation_config)
|
||
# 获取数据列表
|
||
with open(data_list_path, 'r') as f:
|
||
self.lines = f.readlines()
|
||
|
||
def __getitem__(self, idx):
|
||
# 分割音频路径和标签
|
||
audio_path, label = self.lines[idx].replace('\n', '').split('\t')
|
||
# 读取音频
|
||
audio_segment = AudioSegment.from_file(audio_path)
|
||
# 裁剪静音
|
||
if self.do_vad:
|
||
audio_segment.vad()
|
||
# 数据太短不利于训练
|
||
if self.mode == 'train':
|
||
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)
|
||
# decibel normalization
|
||
if self._use_dB_normalization:
|
||
audio_segment.normalize(target_db=self._target_dB)
|
||
# 裁剪需要的数据
|
||
audio_segment.crop(duration=self.max_duration, mode=self.mode)
|
||
# 音频增强
|
||
self._augmentation_pipeline.transform_audio(audio_segment)
|
||
return np.array(audio_segment.samples, dtype=np.float32), np.array(int(label), dtype=np.int64)
|
||
|
||
def __len__(self):
|
||
return len(self.lines)
|