import os import torch from torch.utils.data import Dataset, DataLoader, random_split, Subset import torchaudio from sklearn.model_selection import KFold import glob import numpy as np from utils.utils import repeat_expand # def get_file_name(path): # normalized_path = os.path.normpath(path) # path_parts = normalized_path.split(os.sep) # spk_name = path_parts[-2] # wav_name = path_parts[-1] # file_name = f'{spk_name}_{wav_name}' # return file_name # # Define the dataset # class AudioDataset(Dataset): # def __init__(self, audio_paths, transform=None): # """ # 初始化数据集。 # :param audio_paths: 音频文件的路径列表。 # :param transform: 应用于每个音频样本的可选变换。 # """ # self.audio_paths = audio_paths # self.transform = transform # def __len__(self): # """ # 返回数据集中样本的数量。 # """ # return len(self.audio_paths) # def __getitem__(self, idx): # """ # 根据索引获取音频样本。 # """ # audio_16k_path = self.audio_paths[idx] # mel_44k_path = audio_16k_path.replace('audio_16k', 'mel_44k').replace('.wav', '.npy') # # 加载音频文件 # # wav, sr = torchaudio.load(audio_path) # wav_16k ,_ = torchaudio.load(audio_16k_path) # mel_44k = torch.from_numpy(np.load(mel_44k_path)) # file_name = get_file_name(audio_16k_path) # return wav_16k, mel_44k, file_name def get_file_name(path): normalized_path = os.path.normpath(path) path_parts = normalized_path.split(os.sep) try: # 尝试获取倒数第二个和最后一个部分作为说话者名和文件名 spk_name = path_parts[-2] wav_name = path_parts[-1].split('.')[0] # 移除文件扩展名 file_name = f'{spk_name}_{wav_name}' except IndexError: # 如果路径格式不正确，返回原始路径 file_name = path return file_name def wav_pad(wav, multiple=200): batch, seq_len = wav.shape padded_len = ((seq_len + (multiple-1)) // multiple) * multiple padded_wav = repeat_expand(wav, padded_len) return padded_wav class AudioDataset(torch.utils.data.Dataset): def __init__(self, audio_paths, transform=None): """ 初始化数据集。 :param audio_paths: 音频文件的路径列表。 :param transform: 应用于每个音频样本的可选变换。 """ self.audio_paths = audio_paths self.transform = transform def __len__(self): """ 返回数据集中样本的数量。 """ return len(self.audio_paths) def __getitem__(self, idx): """ 根据索引获取音频样本。 """ if idx >= len(self.audio_paths): raise IndexError("Index out of bounds") audio_16k_path = self.audio_paths[idx] mel_44k_path = os.path.splitext(audio_16k_path)[0].replace('audio_16k', 'mel_44k') + '.npy' audio_44k_path = os.path.splitext(audio_16k_path)[0].replace('audio_16k', 'audio_44k') + '.wav' vq_post_path = os.path.splitext(audio_16k_path)[0].replace('audio_16k', 'vq_post') + '.npy' spk_path = os.path.splitext(audio_16k_path)[0].replace('audio_16k', 'spk') + '.npy' wav_44k, _ = torchaudio.load(audio_44k_path) # For padding to muliple of 200 wav_44k = wav_pad(wav_44k) # 确保mel_44k文件存在 if not os.path.isfile(mel_44k_path): raise FileNotFoundError(f"Mel spectrogram file not found: {mel_44k_path}") mel_44k = torch.from_numpy(np.load(mel_44k_path)) vq_post = torch.from_numpy(np.load(vq_post_path)) spk = torch.from_numpy(np.load(spk_path)) file_name = get_file_name(audio_16k_path) return wav_44k, mel_44k, vq_post, spk, file_name def load_audio_data(data_path, batch_size=64, validation_ratio=0.1, demo_num=0): # train_audio_paths = [os.path.join(data_path, 'train', f) for f in os.listdir(data_path) if f.endswith('.wav')] # test_audio_paths = [os.path.join(data_path, 'test', f) for f in os.listdir(data_path) if f.endswith('.wav')] train_audio_16k_paths = glob.glob(os.path.join(data_path, 'train', 'audio_16k/**', '*.wav')) test_audio_16k_paths = glob.glob(os.path.join(data_path, 'test', 'audio_16k/**', '*.wav')) if type(demo_num) is int and demo_num > 0: try: train_audio_16k_paths = train_audio_16k_paths[:demo_num] test_audio_16k_paths = test_audio_16k_paths[:demo_num//9] except: raise ValueError dataset = AudioDataset(train_audio_16k_paths) # Split dataset into training and validation sets total_size = len(dataset) val_size = int(validation_ratio * total_size) train_size = total_size - val_size train_dataset, val_dataset = random_split(dataset, [train_size, val_size]) test_dataset = AudioDataset(test_audio_16k_paths) # Create data loaders train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True) val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=True) test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False) return train_loader, val_loader, test_loader def load_audio_data_k_fold(data_path, k_folds=5, batch_size=64, demo_num=0): audio_paths = [os.path.join(data_path, f) for f in os.listdir(data_path) if f.endswith('.wav')] if type(demo_num) is int and demo_num > 0: try: audio_paths = audio_paths[:demo_num] except: raise ValueError # 创建一个完整的数据集 full_dataset = AudioDataset(audio_paths) # 初始化KFold kf = KFold(n_splits=k_folds, shuffle=True, random_state=42) fold_loaders = [] for train_index, val_index in kf.split(full_dataset): # 根据索引分割数据集 train_dataset = Subset(full_dataset, train_index) val_dataset = Subset(full_dataset, val_index) # 创建训练和验证的数据加载器 train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True) val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=False) # 添加当前折的数据加载器到列表 fold_loaders.append((train_loader, val_loader)) # 测试集数据加载器不需要分割，使用全部数据 test_loader = DataLoader(dataset=full_dataset, batch_size=batch_size, shuffle=False) return fold_loaders, test_loader class FocalLoss: def __init__(self, alpha_t=None, gamma=0): """ :param alpha_t: A list of weights for each class :param gamma: """ self.alpha_t = torch.tensor(alpha_t) if alpha_t else None self.gamma = gamma def __call__(self, outputs, targets): if self.alpha_t is None and self.gamma == 0: focal_loss = torch.nn.functional.cross_entropy(outputs, targets) elif self.alpha_t is not None and self.gamma == 0: if self.alpha_t.device != outputs.device: self.alpha_t = self.alpha_t.to(outputs) focal_loss = torch.nn.functional.cross_entropy(outputs, targets, weight=self.alpha_t) elif self.alpha_t is None and self.gamma != 0: ce_loss = torch.nn.functional.cross_entropy(outputs, targets, reduction='none') p_t = torch.exp(-ce_loss) focal_loss = ((1 - p_t) ** self.gamma * ce_loss).mean() elif self.alpha_t is not None and self.gamma != 0: if self.alpha_t.device != outputs.device: self.alpha_t = self.alpha_t.to(outputs) ce_loss = torch.nn.functional.cross_entropy(outputs, targets, reduction='none') p_t = torch.exp(-ce_loss) ce_loss = torch.nn.functional.cross_entropy(outputs, targets, weight=self.alpha_t, reduction='none') focal_loss = ((1 - p_t) ** self.gamma * ce_loss).mean() # mean over the batch return focal_loss