src/data/dataset.py

import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
import os
import random
from typing import Dict, List, Optional, Tuple
from .audio_processor import AudioProcessor
from ..configs.config import AudioConfig, Config

class SpeakerDataset(Dataset):
    """
    说话人数据集：用于加载单个说话人的音频数据
    """
    def __init__(
        self,
        audio_files: List[str],
        audio_processor: AudioProcessor,
        cache_size: int = 100  # 添加缓存机制
    ):
        self.audio_files = audio_files
        self.audio_processor = audio_processor
        self.cache = {}
        self.cache_size = cache_size

    def __len__(self) -> int:
        return len(self.audio_files)

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        audio_path = self.audio_files[idx]
        
        # 使用缓存机制
        if audio_path in self.cache:
            return self.cache[audio_path]
            
        try:
            audio, mel_spec = self.audio_processor.preprocess_audio(audio_path)
            item = {
                'audio': torch.FloatTensor(audio),
                'mel_spec': torch.FloatTensor(mel_spec),
                'file_path': audio_path
            }
            
            # 更新缓存
            if len(self.cache) < self.cache_size:
                self.cache[audio_path] = item
                
            return item
        except Exception as e:
            print(f"Error processing file {audio_path}: {str(e)}")
            # 返回数据集中的下一个有效样本
            return self.__getitem__((idx + 1) % len(self))

class VoiceDatasetManager:
    """
    数据集管理器：负责数据集的组织和任务采样
    """
    def __init__(
        self,
        root_dir: str,
        audio_processor: Optional[AudioProcessor] = None,
        config: Optional[Config] = None
    ):
        self.root_dir = root_dir
        self.config = config or Config()
        self.audio_processor = audio_processor or AudioProcessor(config=self.config.audio)
        self.speakers = self._scan_speakers()
        
    def _scan_speakers(self) -> Dict[str, List[str]]:
        speakers = {}
        for speaker_id in os.listdir(self.root_dir):
            speaker_dir = os.path.join(self.root_dir, speaker_id)
            if os.path.isdir(speaker_dir):
                audio_files = []
                # 递归搜索所有子目录
                for root, _, files in os.walk(speaker_dir):
                    for file in files:
                        if file.endswith(self.config.data.valid_audio_extensions):
                            audio_path = os.path.join(root, file)
                            # 验证文件是否可访问
                            if os.path.exists(audio_path) and os.path.getsize(audio_path) > 0:
                                audio_files.append(audio_path)
                                
                # 只保留具有足够样本的说话人
                if len(audio_files) >= self.config.data.min_samples_per_speaker:
                    speakers[speaker_id] = audio_files
                else:
                    print(f"Warning: Speaker {speaker_id} has insufficient samples")
        return speakers

    def get_speaker_dataset(self, speaker_id: str) -> SpeakerDataset:
        """获取特定说话人的数据集"""
        if speaker_id not in self.speakers:
            raise ValueError(f"Speaker {speaker_id} not found in dataset")
        return SpeakerDataset(
            self.speakers[speaker_id],
            self.audio_processor,
            cache_size=self.config.data.cache_size
        )

class MetaLearningDataset(Dataset):
    """
    元学习数据集：用于少样本语音克隆的训练
    每次返回一个任务的数据，包含支持集和查询集
    """
    def __init__(
        self,
        dataset_manager: VoiceDatasetManager,
        config: Config
    ):
        self.dataset_manager = dataset_manager
        self.config = config
        
        # 验证数据集
        available_speakers = [
            speaker_id for speaker_id, files in dataset_manager.speakers.items()
            if len(files) >= (config.meta_learning.k_shot + config.meta_learning.k_query)
        ]
        
        if len(available_speakers) < config.meta_learning.n_way:
            raise ValueError(
                f"Not enough speakers with sufficient samples. "
                f"Need {config.meta_learning.n_way} speakers with "
                f"{config.meta_learning.k_shot + config.meta_learning.k_query} samples each, "
                f"but only found {len(available_speakers)}"
            )
            
        self.available_speakers = available_speakers

    def __len__(self) -> int:
        return self.config.meta_learning.n_tasks

    def __getitem__(self, idx: int) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
        """
        返回一个任务的数据
        Returns:
            support_data: 包含支持集数据的字典
                - mel_spec: [n_way*k_shot, n_mels, time]
                - speaker_ids: [n_way*k_shot]
            query_data: 包含查询集数据的字典
                - mel_spec: [n_way*k_query, n_mels, time]
                - speaker_ids: [n_way*k_query]
        """
        # 随机选择说话人
        selected_speakers = random.sample(self.available_speakers, self.config.meta_learning.n_way)
        
        support_data = {
            'mel_spec': [],
            'speaker_ids': []
        }
        
        query_data = {
            'mel_spec': [],
            'speaker_ids': []
        }
        
        for speaker_idx, speaker_id in enumerate(selected_speakers):
            speaker_files = self.dataset_manager.speakers[speaker_id]
            selected_files = random.sample(
                speaker_files,
                self.config.meta_learning.k_shot + self.config.meta_learning.k_query
            )
            
            for i, file_path in enumerate(selected_files):
                try:
                    _, mel_spec = self.dataset_manager.audio_processor.preprocess_audio(file_path)
                    mel_tensor = torch.FloatTensor(mel_spec)  # [n_mels, time]
                    
                    target_dict = support_data if i < self.config.meta_learning.k_shot else query_data
                    target_dict['mel_spec'].append(mel_tensor)
                    target_dict['speaker_ids'].append(speaker_idx)
                    
                except Exception as e:
                    print(f"Error processing {file_path}: {str(e)}")
                    continue
        
        # 转换为张量
        for data_dict in [support_data, query_data]:
            if len(data_dict['mel_spec']) == 0:
                raise RuntimeError("No valid samples found for task")
                
            data_dict['mel_spec'] = torch.stack(data_dict['mel_spec'])
            data_dict['speaker_ids'] = torch.LongTensor(data_dict['speaker_ids'])
            
        return support_data, query_data

def create_meta_learning_dataloader(
    root_dir: str,
    config: Optional[Config] = None,
    **kwargs
) -> DataLoader:
    """
    创建用于元学习的数据加载器
    
    Args:
        root_dir: 数据集根目录
        config: 配置对象
        **kwargs: 其他参数
        
    Returns:
        DataLoader: 元学习数据加载器
    """
    config = config or Config()
    
    # 更新配置
    for key, value in kwargs.items():
        if hasattr(config.meta_learning, key):
            setattr(config.meta_learning, key, value)
    
    # 创建数据集管理器
    dataset_manager = VoiceDatasetManager(root_dir, config=config)
    
    # 创建数据集
    dataset = MetaLearningDataset(dataset_manager, config)
    
    # 创建数据加载器
    return DataLoader(
        dataset,
        batch_size=1,  # 固定为1，因为每个样本已经包含了一个完整的任务
        shuffle=True,
        num_workers=0,  # 避免多进程带来的问题
        pin_memory=True,
        collate_fn=lambda x: x[0]  # 移除批次维度
    )