""" Data pipeline management and preprocessing for ML trainers. This module provides comprehensive data handling including audio preprocessing, data augmentation, dataset creation, and batch loading with support for various audio formats and training requirements. """ import os import logging import random import numpy as np from pathlib import Path from typing import Dict, List, Tuple, Optional, Union, Callable, Any from dataclasses import dataclass import warnings import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader, random_split, WeightedRandomSampler import torchaudio import torchaudio.transforms as T from sklearn.model_selection import train_test_split # Suppress torchaudio warnings warnings.filterwarnings("ignore", category=UserWarning, module="torchaudio") @dataclass class AudioProcessingConfig: """Configuration for audio processing.""" sample_rate: int = 16000 target_length: float = 1.5 # seconds n_mels: int = 40 n_fft: int = 512 hop_length: int = 160 # 10ms at 16kHz win_length: int = 400 # 25ms at 16kHz fmin: float = 80.0 fmax: float = 8000.0 power: float = 2.0 normalized: bool = True feature_type: str = "mel_spectrogram" # mel_spectrogram, mfcc, raw @property def target_samples(self) -> int: """Get target number of samples.""" return int(self.sample_rate * self.target_length) class AudioProcessor: """Audio processing utilities for the training pipeline.""" def __init__(self, config: AudioProcessingConfig, device: str = "cpu"): """ Initialize audio processor. Args: config: Audio processing configuration device: Device for processing ('cpu', 'cuda') """ self.config = config self.device = device # Initialize transforms self.mel_spectrogram = T.MelSpectrogram( sample_rate=config.sample_rate, n_fft=config.n_fft, hop_length=config.hop_length, win_length=config.win_length, n_mels=config.n_mels, fmin=config.fmin, fmax=config.fmax, power=config.power, normalized=config.normalized ).to(device) self.mfcc = T.MFCC( sample_rate=config.sample_rate, n_mfcc=config.n_mels, melkwargs={ 'n_fft': config.n_fft, 'hop_length': config.hop_length, 'win_length': config.win_length, 'n_mels': config.n_mels, 'fmin': config.fmin, 'fmax': config.fmax } ).to(device) self.amplitude_to_db = T.AmplitudeToDB().to(device) def load_audio(self, filepath: str) -> Tuple[torch.Tensor, int]: """ Load audio file and return waveform and sample rate. Args: filepath: Path to audio file Returns: Tuple of (waveform, sample_rate) """ try: waveform, sr = torchaudio.load(filepath) return waveform, sr except Exception as e: raise ValueError(f"Error loading audio file {filepath}: {str(e)}") def resample(self, waveform: torch.Tensor, orig_sr: int) -> torch.Tensor: """ Resample audio to target sample rate. Args: waveform: Input waveform orig_sr: Original sample rate Returns: Resampled waveform """ if orig_sr != self.config.sample_rate: resampler = T.Resample(orig_sr, self.config.sample_rate) waveform = resampler(waveform) return waveform def normalize_length(self, waveform: torch.Tensor, method: str = "pad_trim") -> torch.Tensor: """ Normalize audio length to target duration. Args: waveform: Input waveform method: Method for length normalization ('pad_trim', 'stretch', 'random_crop') Returns: Length-normalized waveform """ target_samples = self.config.target_samples current_samples = waveform.shape[-1] if method == "pad_trim": if current_samples > target_samples: # Trim from center start = (current_samples - target_samples) // 2 waveform = waveform[..., start:start + target_samples] elif current_samples < target_samples: # Pad with zeros pad_amount = target_samples - current_samples pad_left = pad_amount // 2 pad_right = pad_amount - pad_left waveform = F.pad(waveform, (pad_left, pad_right)) elif method == "random_crop": if current_samples > target_samples: # Random crop start = random.randint(0, current_samples - target_samples) waveform = waveform[..., start:start + target_samples] elif current_samples < target_samples: # Pad with zeros pad_amount = target_samples - current_samples waveform = F.pad(waveform, (0, pad_amount)) elif method == "stretch": # Time-stretch to exact length (may change pitch) if current_samples != target_samples: stretch_factor = target_samples / current_samples # Simple linear interpolation for stretching indices = torch.linspace(0, current_samples - 1, target_samples) waveform = torch.nn.functional.interpolate( waveform.unsqueeze(0), size=target_samples, mode='linear', align_corners=True ).squeeze(0) return waveform def extract_features(self, waveform: torch.Tensor) -> torch.Tensor: """ Extract features from waveform. Args: waveform: Input waveform Returns: Feature tensor """ waveform = waveform.to(self.device) if self.config.feature_type == "mel_spectrogram": features = self.mel_spectrogram(waveform) features = self.amplitude_to_db(features) elif self.config.feature_type == "mfcc": features = self.mfcc(waveform) elif self.config.feature_type == "raw": features = waveform else: raise ValueError(f"Unknown feature type: {self.config.feature_type}") return features def process_audio(self, filepath: str, normalize_method: str = "pad_trim") -> torch.Tensor: """ Complete audio processing pipeline. Args: filepath: Path to audio file normalize_method: Method for length normalization Returns: Processed feature tensor """ # Load audio waveform, sr = self.load_audio(filepath) # Convert to mono if stereo if waveform.shape[0] > 1: waveform = torch.mean(waveform, dim=0, keepdim=True) # Resample waveform = self.resample(waveform, sr) # Normalize length waveform = self.normalize_length(waveform, normalize_method) # Extract features features = self.extract_features(waveform) return features def trim_silence(self, waveform: torch.Tensor, threshold: float = 0.01) -> torch.Tensor: """ Trim silence from beginning and end of audio. Args: waveform: Input waveform threshold: Silence threshold (relative to max amplitude) Returns: Trimmed waveform """ # Calculate energy energy = torch.abs(waveform) max_energy = torch.max(energy) # Find non-silent regions non_silent = energy > (threshold * max_energy) if torch.any(non_silent): # Find first and last non-silent samples non_silent_indices = torch.where(non_silent.squeeze())[0] start_idx = non_silent_indices[0].item() end_idx = non_silent_indices[-1].item() + 1 return waveform[..., start_idx:end_idx] else: # If all silent, return original return waveform class AudioAugmentation: """Audio data augmentation utilities.""" def __init__(self, sample_rate: int = 16000, device: str = "cpu"): """Initialize audio augmentation.""" self.sample_rate = sample_rate self.device = device def add_noise(self, waveform: torch.Tensor, noise_factor: float = 0.1) -> torch.Tensor: """Add random noise to audio.""" noise = torch.randn_like(waveform) * noise_factor return waveform + noise def time_shift(self, waveform: torch.Tensor, shift_factor: float = 0.1) -> torch.Tensor: """Apply random time shift.""" shift_amount = int(waveform.shape[-1] * shift_factor * (2 * random.random() - 1)) return torch.roll(waveform, shift_amount, dims=-1) def speed_change(self, waveform: torch.Tensor, speed_factor: float = 0.1) -> torch.Tensor: """Change playback speed (affects both tempo and pitch).""" speed_change = 1.0 + speed_factor * (2 * random.random() - 1) # Simple resampling for speed change new_length = int(waveform.shape[-1] / speed_change) indices = torch.linspace(0, waveform.shape[-1] - 1, new_length) # Linear interpolation waveform_resampled = torch.nn.functional.interpolate( waveform.unsqueeze(0), size=new_length, mode='linear', align_corners=True ).squeeze(0) return waveform_resampled def pitch_shift(self, waveform: torch.Tensor, pitch_factor: float = 0.1) -> torch.Tensor: """Apply pitch shift (simplified version).""" # This is a simplified pitch shift - in practice, you'd use more sophisticated methods pitch_change = 1.0 + pitch_factor * (2 * random.random() - 1) # Create a simple frequency domain shift fft = torch.fft.fft(waveform) shifted_fft = torch.zeros_like(fft) n = fft.shape[-1] shift_bins = int(n * (pitch_change - 1.0)) if shift_bins > 0: shifted_fft[..., shift_bins:] = fft[..., :-shift_bins] elif shift_bins < 0: shifted_fft[..., :shift_bins] = fft[..., -shift_bins:] else: shifted_fft = fft return torch.fft.ifft(shifted_fft).real def volume_change(self, waveform: torch.Tensor, volume_factor: float = 0.2) -> torch.Tensor: """Apply random volume change.""" volume_change = 1.0 + volume_factor * (2 * random.random() - 1) return waveform * volume_change def apply_augmentations(self, waveform: torch.Tensor, augmentation_config: Dict[str, float]) -> torch.Tensor: """Apply multiple augmentations based on configuration.""" augmented = waveform.clone() if augmentation_config.get('noise_factor', 0) > 0: if random.random() < 0.5: # 50% chance augmented = self.add_noise(augmented, augmentation_config['noise_factor']) if augmentation_config.get('speed_factor', 0) > 0: if random.random() < 0.3: # 30% chance augmented = self.speed_change(augmented, augmentation_config['speed_factor']) if augmentation_config.get('pitch_factor', 0) > 0: if random.random() < 0.3: # 30% chance augmented = self.pitch_shift(augmented, augmentation_config['pitch_factor']) if augmentation_config.get('volume_factor', 0) > 0: if random.random() < 0.4: # 40% chance augmented = self.volume_change(augmented, augmentation_config['volume_factor']) if augmentation_config.get('time_shift_factor', 0) > 0: if random.random() < 0.3: # 30% chance augmented = self.time_shift(augmented, augmentation_config['time_shift_factor']) return augmented class AudioDataset(Dataset): """Generic audio dataset for training.""" def __init__(self, audio_files: List[str], labels: List[int], audio_processor: AudioProcessor, augmentation: Optional[AudioAugmentation] = None, augmentation_config: Optional[Dict[str, float]] = None, apply_augmentation: bool = True): """ Initialize audio dataset. Args: audio_files: List of audio file paths labels: List of corresponding labels audio_processor: Audio processor instance augmentation: Audio augmentation instance augmentation_config: Augmentation configuration apply_augmentation: Whether to apply augmentation during training """ self.audio_files = audio_files self.labels = labels self.audio_processor = audio_processor self.augmentation = augmentation self.augmentation_config = augmentation_config or {} self.apply_augmentation = apply_augmentation assert len(audio_files) == len(labels), "Audio files and labels must have same length" def __len__(self) -> int: return len(self.audio_files) def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]: """Get item from dataset.""" audio_file = self.audio_files[idx] label = self.labels[idx] try: # Load and process audio waveform, sr = self.audio_processor.load_audio(audio_file) # Convert to mono if waveform.shape[0] > 1: waveform = torch.mean(waveform, dim=0, keepdim=True) # Resample waveform = self.audio_processor.resample(waveform, sr) # Apply augmentation if enabled and available if (self.apply_augmentation and self.augmentation is not None and self.augmentation_config): waveform = self.augmentation.apply_augmentations( waveform, self.augmentation_config ) # Normalize length and extract features waveform = self.audio_processor.normalize_length(waveform) features = self.audio_processor.extract_features(waveform) # Ensure consistent shape if features.dim() == 3: # (1, freq, time) features = features.squeeze(0) # (freq, time) return features, label except Exception as e: # Return zero tensor and label on error logging.warning(f"Error processing {audio_file}: {str(e)}") dummy_features = torch.zeros( self.audio_processor.config.n_mels, self.audio_processor.config.target_samples // self.audio_processor.config.hop_length + 1 ) return dummy_features, label class DatasetBuilder: """Builder for creating training datasets from file structures.""" def __init__(self, logger: Optional[logging.Logger] = None): """Initialize dataset builder.""" self.logger = logger or logging.getLogger(__name__) def scan_directory_structure(self, data_dir: str, supported_extensions: List[str] = None) -> Dict[str, List[str]]: """ Scan directory structure and organize files by class. Args: data_dir: Root data directory supported_extensions: List of supported file extensions Returns: Dictionary mapping class names to file lists """ if supported_extensions is None: supported_extensions = ['.wav', '.mp3', '.flac', '.m4a'] data_dir = Path(data_dir) if not data_dir.exists(): raise FileNotFoundError(f"Data directory not found: {data_dir}") class_files = {} # For each subdirectory (class) for class_dir in data_dir.iterdir(): if class_dir.is_dir(): class_name = class_dir.name files = [] # Find all audio files in class directory for ext in supported_extensions: files.extend(list(class_dir.glob(f"**/*{ext}"))) if files: class_files[class_name] = [str(f) for f in files] self.logger.info(f"Found {len(files)} files for class '{class_name}'") return class_files def create_wakeword_dataset(self, data_dir: str, audio_processor: AudioProcessor, wakeword_classes: Dict[int, str], val_split: float = 0.2, test_split: float = 0.1, augmentation_config: Dict[str, float] = None) -> Tuple[AudioDataset, AudioDataset, AudioDataset]: """ Create wakeword detection dataset. Args: data_dir: Data directory containing wakeword audio files audio_processor: Audio processor instance wakeword_classes: Mapping of class IDs to class names val_split: Validation split ratio test_split: Test split ratio augmentation_config: Augmentation configuration Returns: Tuple of (train_dataset, val_dataset, test_dataset) """ # Scan directory structure class_files = self.scan_directory_structure(data_dir) # Map class names to IDs class_name_to_id = {name: class_id for class_id, name in wakeword_classes.items()} all_files = [] all_labels = [] for class_name, files in class_files.items(): if class_name in class_name_to_id: class_id = class_name_to_id[class_name] all_files.extend(files) all_labels.extend([class_id] * len(files)) self.logger.info(f"Added {len(files)} files for class '{class_name}' (ID: {class_id})") else: self.logger.warning(f"Unknown class '{class_name}' - skipping") if not all_files: raise ValueError("No valid audio files found") # Split data train_files, test_files, train_labels, test_labels = train_test_split( all_files, all_labels, test_size=test_split, random_state=42, stratify=all_labels ) train_files, val_files, train_labels, val_labels = train_test_split( train_files, train_labels, test_size=val_split/(1-test_split), random_state=42, stratify=train_labels ) # Create augmentation augmentation = AudioAugmentation(audio_processor.config.sample_rate) # Create datasets train_dataset = AudioDataset( train_files, train_labels, audio_processor, augmentation, augmentation_config, apply_augmentation=True ) val_dataset = AudioDataset( val_files, val_labels, audio_processor, None, None, apply_augmentation=False ) test_dataset = AudioDataset( test_files, test_labels, audio_processor, None, None, apply_augmentation=False ) self.logger.info(f"Created datasets - Train: {len(train_dataset)}, " f"Val: {len(val_dataset)}, Test: {len(test_dataset)}") return train_dataset, val_dataset, test_dataset def create_voice_recognition_dataset(self, data_dir: str, audio_processor: AudioProcessor, speaker_mapping: Dict[str, str], val_split: float = 0.2, test_split: float = 0.1, augmentation_config: Dict[str, float] = None) -> Tuple[AudioDataset, AudioDataset, AudioDataset, Dict[str, str]]: """ Create voice recognition dataset. Args: data_dir: Data directory containing speaker audio files audio_processor: Audio processor instance speaker_mapping: Mapping of directory names to speaker names val_split: Validation split ratio test_split: Test split ratio augmentation_config: Augmentation configuration Returns: Tuple of (train_dataset, val_dataset, test_dataset, final_speaker_mapping) """ # Scan directory structure class_files = self.scan_directory_structure(data_dir) # Create speaker ID mapping speaker_id_mapping = {} all_files = [] all_labels = [] for idx, (dir_name, files) in enumerate(class_files.items()): speaker_name = speaker_mapping.get(dir_name, dir_name) speaker_id = f"speaker{str(idx).zfill(3)}" speaker_id_mapping[speaker_id] = speaker_name all_files.extend(files) all_labels.extend([idx] * len(files)) self.logger.info(f"Added {len(files)} files for speaker '{speaker_name}' " f"(ID: {speaker_id}, Label: {idx})") if not all_files: raise ValueError("No valid audio files found") # Split data train_files, test_files, train_labels, test_labels = train_test_split( all_files, all_labels, test_size=test_split, random_state=42, stratify=all_labels ) train_files, val_files, train_labels, val_labels = train_test_split( train_files, train_labels, test_size=val_split/(1-test_split), random_state=42, stratify=train_labels ) # Create augmentation augmentation = AudioAugmentation(audio_processor.config.sample_rate) # Create datasets train_dataset = AudioDataset( train_files, train_labels, audio_processor, augmentation, augmentation_config, apply_augmentation=True ) val_dataset = AudioDataset( val_files, val_labels, audio_processor, None, None, apply_augmentation=False ) test_dataset = AudioDataset( test_files, test_labels, audio_processor, None, None, apply_augmentation=False ) self.logger.info(f"Created voice recognition datasets - Train: {len(train_dataset)}, " f"Val: {len(val_dataset)}, Test: {len(test_dataset)}") return train_dataset, val_dataset, test_dataset, speaker_id_mapping class DataLoaderFactory: """Factory for creating optimized data loaders.""" @staticmethod def create_dataloader(dataset: Dataset, batch_size: int, shuffle: bool = True, num_workers: int = 4, pin_memory: bool = True, drop_last: bool = False, balanced_sampling: bool = False) -> DataLoader: """ Create optimized data loader. Args: dataset: Dataset to load from batch_size: Batch size shuffle: Whether to shuffle data num_workers: Number of worker processes pin_memory: Whether to pin memory drop_last: Whether to drop last incomplete batch balanced_sampling: Whether to use balanced sampling for imbalanced datasets Returns: Configured DataLoader """ sampler = None if balanced_sampling and hasattr(dataset, 'labels'): # Create weighted sampler for balanced sampling labels = dataset.labels class_counts = torch.bincount(torch.tensor(labels)) class_weights = 1.0 / class_counts.float() sample_weights = class_weights[labels] sampler = WeightedRandomSampler( weights=sample_weights, num_samples=len(dataset), replacement=True ) shuffle = False # Don't shuffle when using sampler return DataLoader( dataset=dataset, batch_size=batch_size, shuffle=shuffle, sampler=sampler, num_workers=num_workers, pin_memory=pin_memory, drop_last=drop_last, persistent_workers=num_workers > 0 ) class DataPreprocessor: """Comprehensive data preprocessing pipeline.""" def __init__(self, audio_config: AudioProcessingConfig, logger: Optional[logging.Logger] = None): """Initialize data preprocessor.""" self.audio_config = audio_config self.logger = logger or logging.getLogger(__name__) self.audio_processor = AudioProcessor(audio_config) self.dataset_builder = DatasetBuilder(logger) def preprocess_wakeword_data(self, raw_data_dir: str, output_dir: str, wakeword_classes: Dict[int, str]) -> str: """ Preprocess raw wakeword data. This method handles the requirements from the project description: 1. Trim audio files to consistent length 2. Split negative/background files into chunks 3. Sort into train/val/test directories Args: raw_data_dir: Directory containing raw audio files output_dir: Output directory for processed data wakeword_classes: Mapping of class IDs to names Returns: Path to processed data directory """ raw_data_path = Path(raw_data_dir) output_path = Path(output_dir) self.logger.info(f"Preprocessing wakeword data from {raw_data_dir}") # Process each class for class_id, class_name in wakeword_classes.items(): class_raw_dir = raw_data_path / class_name if not class_raw_dir.exists(): self.logger.warning(f"Class directory not found: {class_raw_dir}") continue self.logger.info(f"Processing class: {class_name}") # Get all audio files audio_files = [] for ext in ['.wav', '.mp3', '.flac']: audio_files.extend(list(class_raw_dir.glob(f"**/*{ext}"))) if not audio_files: self.logger.warning(f"No audio files found for class: {class_name}") continue processed_files = [] for audio_file in audio_files: try: # Load and process audio waveform, sr = self.audio_processor.load_audio(str(audio_file)) # Convert to mono if waveform.shape[0] > 1: waveform = torch.mean(waveform, dim=0, keepdim=True) # Resample waveform = self.audio_processor.resample(waveform, sr) if class_name in ['negative', 'background']: # Split long files into chunks chunks = self._split_into_chunks(waveform, self.audio_config.target_samples) for i, chunk in enumerate(chunks): chunk_filename = f"{audio_file.stem}_chunk_{i:03d}.wav" processed_files.append((chunk, chunk_filename)) else: # For wakeword files, trim silence and normalize length waveform = self.audio_processor.trim_silence(waveform) waveform = self.audio_processor.normalize_length(waveform) processed_files.append((waveform, audio_file.name)) except Exception as e: self.logger.error(f"Error processing {audio_file}: {str(e)}") continue # Save processed files class_output_dir = output_path / "chunked" / class_name class_output_dir.mkdir(parents=True, exist_ok=True) for waveform, filename in processed_files: output_file = class_output_dir / filename torchaudio.save(str(output_file), waveform, self.audio_config.sample_rate) self.logger.info(f"Processed {len(processed_files)} files for class: {class_name}") # Split into train/val/test self._split_processed_data(output_path / "chunked", output_path) return str(output_path) def preprocess_voice_recognition_data(self, raw_data_dir: str, output_dir: str, target_chunk_length: float = None) -> str: """ Preprocess raw voice recognition data. Args: raw_data_dir: Directory containing raw speaker audio files output_dir: Output directory for processed data target_chunk_length: Target chunk length (uses wakeword length if None) Returns: Path to processed data directory """ if target_chunk_length is None: target_chunk_length = self.audio_config.target_length raw_data_path = Path(raw_data_dir) output_path = Path(output_dir) self.logger.info(f"Preprocessing voice recognition data from {raw_data_dir}") # Process each speaker directory for speaker_dir in raw_data_path.iterdir(): if not speaker_dir.is_dir(): continue speaker_name = speaker_dir.name self.logger.info(f"Processing speaker: {speaker_name}") # Get all audio files audio_files = [] for ext in ['.wav', '.mp3', '.flac']: audio_files.extend(list(speaker_dir.glob(f"**/*{ext}"))) if not audio_files: self.logger.warning(f"No audio files found for speaker: {speaker_name}") continue processed_files = [] target_samples = int(self.audio_config.sample_rate * target_chunk_length) for audio_file in audio_files: try: # Load and process audio waveform, sr = self.audio_processor.load_audio(str(audio_file)) # Convert to mono if waveform.shape[0] > 1: waveform = torch.mean(waveform, dim=0, keepdim=True) # Resample waveform = self.audio_processor.resample(waveform, sr) # Trim silence waveform = self.audio_processor.trim_silence(waveform) # Split into chunks chunks = self._split_into_chunks(waveform, target_samples) for i, chunk in enumerate(chunks): chunk_filename = f"{audio_file.stem}_chunk_{i:03d}.wav" processed_files.append((chunk, chunk_filename)) except Exception as e: self.logger.error(f"Error processing {audio_file}: {str(e)}") continue # Save processed files speaker_output_dir = output_path / "chunked" / speaker_name speaker_output_dir.mkdir(parents=True, exist_ok=True) for waveform, filename in processed_files: output_file = speaker_output_dir / filename torchaudio.save(str(output_file), waveform, self.audio_config.sample_rate) self.logger.info(f"Processed {len(processed_files)} files for speaker: {speaker_name}") # Split into train/val/test self._split_processed_data(output_path / "chunked", output_path) return str(output_path) def _split_into_chunks(self, waveform: torch.Tensor, chunk_size: int, overlap: float = 0.0) -> List[torch.Tensor]: """Split waveform into chunks of specified size.""" chunks = [] total_samples = waveform.shape[-1] if total_samples <= chunk_size: # If shorter than chunk size, pad and return padded = F.pad(waveform, (0, chunk_size - total_samples)) chunks.append(padded) else: # Split into overlapping chunks step_size = int(chunk_size * (1 - overlap)) start = 0 while start + chunk_size <= total_samples: chunk = waveform[..., start:start + chunk_size] chunks.append(chunk) start += step_size # Add final chunk if there's remaining audio if start < total_samples: final_chunk = waveform[..., -chunk_size:] chunks.append(final_chunk) return chunks def _split_processed_data(self, chunked_dir: Path, output_dir: Path): """Split processed data into train/val/test directories.""" train_dir = output_dir / "train" val_dir = output_dir / "val" test_dir = output_dir / "test" # Create directories for split_dir in [train_dir, val_dir, test_dir]: split_dir.mkdir(parents=True, exist_ok=True) # Process each class/speaker for class_dir in chunked_dir.iterdir(): if not class_dir.is_dir(): continue class_name = class_dir.name # Get all files files = list(class_dir.glob("*.wav")) if len(files) < 3: self.logger.warning(f"Too few files for class {class_name} - skipping split") continue # Split files random.shuffle(files) total_files = len(files) test_count = max(1, int(total_files * 0.1)) val_count = max(1, int(total_files * 0.2)) train_count = total_files - test_count - val_count test_files = files[:test_count] val_files = files[test_count:test_count + val_count] train_files = files[test_count + val_count:] # Copy files to appropriate directories for split_name, split_files, split_dir in [ ("train", train_files, train_dir), ("val", val_files, val_dir), ("test", test_files, test_dir) ]: class_split_dir = split_dir / class_name class_split_dir.mkdir(exist_ok=True) for file in split_files: import shutil shutil.copy2(file, class_split_dir / file.name) self.logger.info(f"Split {class_name}: {len(train_files)} train, " f"{len(val_files)} val, {len(test_files)} test")