from pathlib import Path import random import logging import numpy as np import librosa import soundfile as sf import json from typing import List, Optional, Dict, Union, Tuple, Any from torch.utils.data import Dataset, Sampler from tqdm import tqdm from data.augment import StemAugmentation, MixtureAugmentation from data.moise_taxonomy import get_banned_other_pairs, get_target_stem_pairs logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) AUDIO_EXTENSIONS = ['.flac', '.mp3', '.wav'] DEFAULT_GAIN_RANGE = (0.5, 1.0) def calculate_rms(audio: np.ndarray) -> float: return np.sqrt(np.mean(audio**2)) def contains_audio_signal(audio: np.ndarray, rms_threshold: float = 0.01) -> bool: if audio is None: return False return calculate_rms(audio) > rms_threshold def fix_length(target: np.ndarray, source: np.ndarray) -> np.ndarray: target_length, source_length = target.shape[-1], source.shape[-1] if target_length < source_length: return np.pad(target, ((0, 0), (0, source_length - target_length)), mode='constant') if target_length > source_length: return target[:, :source_length] return target def fix_length_to_duration(target: np.ndarray, duration: float, sr: int) -> np.ndarray: target_length = target.shape[-1] required_length = int(duration * sr) if target_length < required_length: return np.pad(target, ((0, 0), (0, required_length - target_length)), mode='constant') if target_length > required_length: return target[:, :required_length] return target def get_audio_duration(file_path: Path) -> float: try: return sf.info(file_path).duration except Exception as e: logger.error(f"Error getting duration for {file_path}: {e}") return 0.0 def mix_to_target_snr(target: np.ndarray, noise: np.ndarray, target_snr_db: float) -> Tuple[np.ndarray, float, float]: target_power, noise_power = np.mean(target ** 2), np.mean(noise ** 2) if noise_power < 1e-8: return target.copy(), 1.0, 0.0 if target_power < 1e-8: return noise * 0.001, 0.0, 0.001 target_snr_linear = 10 ** (target_snr_db / 10) noise_scale = np.sqrt(target_power / (noise_power * target_snr_linear)) scaled_noise = noise * noise_scale mixture = target + scaled_noise max_amplitude = np.max(np.abs(mixture)) target_scale = 1.0 if max_amplitude > 1.0: norm_factor = 0.95 / max_amplitude mixture *= norm_factor target_scale = norm_factor noise_scale *= norm_factor return mixture, target_scale, noise_scale class RawStems(Dataset): def __init__( self, target_stem: str, root_directory: Union[str, Path], sr: int = 48000, clip_duration: float = 3.0, snr_range: Tuple[float, float] = (0.0, 10.0), apply_augmentation: bool = True, also_apply_mixture_augmentation: bool = True, also_apply_stem_augmentation: bool = True, rms_threshold: float = -40.0, no_mixture: bool = False, no_mixture_all: bool = False, output_mixture: bool = False, moisesdb: bool = False, random_mixture: bool = False, ) -> None: self.root_directory = Path(root_directory) self.sr = sr self.clip_duration = clip_duration self.snr_range = snr_range self.apply_augmentation = apply_augmentation self.apply_stem_augmentation = apply_augmentation and also_apply_stem_augmentation self.apply_mixture_augmentation = apply_augmentation and also_apply_mixture_augmentation self.rms_threshold = rms_threshold self.no_mixture = no_mixture self.no_mixture_all = no_mixture_all self.output_mixture = output_mixture self.random_mixture = random_mixture if moisesdb: self.target_stems = get_target_stem_pairs(target_stem) self.allowed_others = ["vocals", "bass", "drums", "guitar", "other_plucked", "percussion", "piano", "other_keys", "bowed_strings", "wind", "other"] self.banned_others = [] else: target_stem_parts = target_stem.split("_") target_stem_1 = target_stem_parts[0].strip() target_stem_2 = target_stem_parts[1].strip() if len(target_stem_parts) > 1 else None self.target_stems = [(target_stem_1, target_stem_2)] self.allowed_others = ["Kbs", "Gtr", "Bass", "Voc", "Synth", "Rhy", "Orch"] assert target_stem_1 in self.allowed_others self.banned_others = [] logger.info(f"Scanning '{self.root_directory}' for songs containing stem '{target_stem}'...") self.folders = [] for song_dir in self.root_directory.iterdir(): if song_dir.is_dir(): for (target_stem_1, target_stem_2) in self.target_stems: target_path = song_dir / target_stem_1 if target_stem_2: target_path /= target_stem_2 if target_path.exists() and target_path.is_dir(): self.folders.append(song_dir) break if not self.folders: raise FileNotFoundError(f"No subdirectories in '{self.root_directory}' were found containing the stem path '{target_stem}'. " f"Please check your directory structure.") logger.info(f"Found {len(self.folders)} song folders.") self.audio_files = self._index_audio_files() if not self.audio_files: raise ValueError("No audio files found.") logger.info(f"Indexed {len(self.audio_files)} audio files.") self.activity_masks = self._compute_activity_masks() self._filter_activity_masks() logger.info(f"{len(self.audio_files)} audio files after filtering.") self.stem_augmentation = StemAugmentation() if self.apply_mixture_augmentation: self.mixture_augmentation = MixtureAugmentation() def load_audio(self, file_path: Path, offset: float, duration: float, sr: int, aug: bool) -> np.ndarray: expected_samples = int(sr * duration) audio, _ = librosa.load(file_path, sr=sr, offset=offset, duration=duration, mono=False) if len(audio.shape) == 1: audio = audio.reshape(1, -1) if audio.shape[1] == 0: return np.zeros((2, expected_samples)) if audio.shape[0] == 1: audio = np.vstack([audio, audio]) if audio.shape[1] > expected_samples: audio = audio[:, :expected_samples] elif audio.shape[1] < expected_samples: padding = expected_samples - audio.shape[1] audio = np.pad(audio, ((0, 0), (0, padding)), mode='constant') if aug and self.apply_stem_augmentation: audio = self.stem_augmentation.apply(audio, self.sr) return audio def _compute_activity_masks(self) -> Dict[str, np.ndarray]: rms_analysis_path = self.root_directory / "rms_analysis.jsonl" if not rms_analysis_path.exists(): logger.warning("rms_analysis.jsonl not found. Non-silent selection will be disabled.") return {} logger.info(f"Loading and processing RMS data from {rms_analysis_path}") rms_data = {} with open(rms_analysis_path, 'r', encoding='utf-8') as f: for line in f: try: data = json.loads(line) rms_data[data['filepath']] = np.array(data['rms_db_per_second']) except (json.JSONDecodeError, KeyError): continue logger.info("Computing activity masks for all indexed files...") activity_masks = {} window_size = int(np.ceil(self.clip_duration)) all_indexed_files = set() for song in self.audio_files: all_indexed_files.update(p.relative_to(self.root_directory) for p in song["target_stems"]) all_indexed_files.update(p.relative_to(self.root_directory) for p in song["others"]) for relative_path in tqdm(all_indexed_files, desc="Computing Activity Masks"): path_str = str(relative_path) if path_str in rms_data: rms_values = rms_data[path_str] if len(rms_values) < window_size: activity_masks[path_str] = np.array([False] * len(rms_values)) continue is_loud = rms_values > self.rms_threshold sum_loud = np.convolve(is_loud, np.ones(window_size), 'valid') avg_loud_enough = sum_loud / window_size > 0.8 mask = np.zeros(len(rms_values), dtype=bool) mask[:len(avg_loud_enough)] = avg_loud_enough activity_masks[path_str] = mask else: logger.warning(f"No RMS data found for {path_str}") return activity_masks def _filter_activity_masks(self) -> None: def filter_stem(stem: Path) -> bool: if not self._find_common_valid_start_seconds([stem]): # logger.warning(f"Skipping {stem} due to silence.") return False return True def filter_song(song: Dict[str, List[Path]]) -> bool: if song["target_stems"] and song["others"]: return True # logger.warning(f"Skipping {song} due to empty or invalid audio.") return False for song in self.audio_files: song["target_stems"] = list(filter(filter_stem, song["target_stems"])) song["others"] = list(filter(filter_stem, song["others"])) self.audio_files = list(filter(filter_song, self.audio_files)) def _find_common_valid_start_seconds(self, file_paths: List[Path]) -> List[int]: if not self.activity_masks: return [] common_mask = None min_len = float('inf') masks_to_intersect = [] for file_path in file_paths: path_str = str(file_path.relative_to(self.root_directory)) mask = self.activity_masks.get(path_str) if mask is None: return [] masks_to_intersect.append(mask) min_len = min(min_len, len(mask)) if not masks_to_intersect: return [] final_mask = np.ones(min_len, dtype=bool) for mask in masks_to_intersect: final_mask &= mask[:min_len] return np.where(final_mask)[0].tolist() def _index_audio_files(self) -> List[Dict[str, List[Path]]]: indexed_songs = [] for folder in tqdm(self.folders, desc="Indexing audio files"): song_dict = {"target_stems": [], "others": []} for (target_stem_1, target_stem_2) in self.target_stems: target_folder = folder / target_stem_1 if target_stem_2: target_folder /= target_stem_2 if target_folder.exists(): song_dict["target_stems"].extend(p for p in target_folder.rglob('*') if p.suffix.lower() in AUDIO_EXTENSIONS) for p in folder.rglob('*'): if p.suffix.lower() in AUDIO_EXTENSIONS: try: relative_path = p.relative_to(folder) parts = relative_path.parts if not (len(parts) > 0 and parts[0] in self.allowed_others): # logger.warning(f"Skipping {p} due to unknown stem.") raise ValueError for (target_stem_1, target_stem_2) in self.target_stems + self.banned_others: if len(parts) > 0 and parts[0] == target_stem_1 and (target_stem_2 is None or (len(parts) > 1 and parts[1] == target_stem_2)): raise ValueError song_dict["others"].append(p) except ValueError: continue if song_dict["target_stems"] and song_dict["others"]: indexed_songs.append(song_dict) # else: # logger.warning(f"Skipping {folder} due to empty or invalid audio.") return indexed_songs def load_other_audio_randomly(self, index: int, offset: float, duration: float, sr: int, aug: bool) -> np.ndarray: song_dict = self.audio_files[index] selected = random.choice(song_dict["others"]) valid_starts = self._find_common_valid_start_seconds([selected]) start_second = random.choice(valid_starts) offset = start_second + random.uniform(0, 1.0 - (self.clip_duration % 1.0 or 1.0)) return self.load_audio(selected, offset, duration, sr, aug) def __getitem__(self, index: int) -> Dict[str, Any]: song_dict = self.audio_files[index] for _ in range(100): if self.no_mixture_all: selected_target = random.choice(song_dict["target_stems"] + song_dict["others"]) valid_starts = self._find_common_valid_start_seconds([selected_target]) if not valid_starts: continue start_second = random.choice(valid_starts) offset = start_second + random.uniform(0, 1.0 - (self.clip_duration % 1.0 or 1.0)) target = self.load_audio(selected_target, offset, self.clip_duration, self.sr, False) if not contains_audio_signal(target): continue target_clean = target.copy() target_augmented = self.stem_augmentation.apply(target, self.sr) if self.apply_stem_augmentation else target return { "mixture": np.nan_to_num(target_augmented), "target": np.nan_to_num(target_clean) } if self.no_mixture: num_targets = random.randint(1, min(len(song_dict["target_stems"]), 5)) selected_targets = random.sample(song_dict["target_stems"], num_targets) valid_starts = self._find_common_valid_start_seconds(selected_targets) if not valid_starts: continue start_second = random.choice(valid_starts) offset = start_second + random.uniform(0, 1.0 - (self.clip_duration % 1.0 or 1.0)) target = sum(self.load_audio(p, offset, self.clip_duration, self.sr, False) for p in selected_targets) / num_targets if not contains_audio_signal(target): continue target_clean = target.copy() target_augmented = self.stem_augmentation.apply(target, self.sr) if self.apply_stem_augmentation else target return { "mixture": np.nan_to_num(target_augmented), "target": np.nan_to_num(target_clean) } num_targets = random.randint(1, min(len(song_dict["target_stems"]), 5)) selected_targets = random.sample(song_dict["target_stems"], num_targets) if not self.random_mixture: num_others = random.randint(1, min(len(song_dict["others"]), 10)) selected_others = random.sample(song_dict["others"], num_others) valid_starts = self._find_common_valid_start_seconds(selected_targets + selected_others) else: valid_starts = self._find_common_valid_start_seconds(selected_targets) if valid_starts: start_second = random.choice(valid_starts) offset = start_second + random.uniform(0, 1.0 - (self.clip_duration % 1.0 or 1.0)) target_mix = sum(self.load_audio(p, offset, self.clip_duration, self.sr, False) for p in selected_targets) / num_targets # aug later if not self.random_mixture: other_mix = sum(self.load_audio(p, offset, self.clip_duration, self.sr, True) for p in selected_others) / num_others else: num_others = random.randint(1, 10) selected_indices = random.sample(range(len(self.audio_files)), num_others) other_mix = sum(self.load_other_audio_randomly(index, offset, self.clip_duration, self.sr, True) for index in selected_indices) / num_others if not contains_audio_signal(target_mix) or not contains_audio_signal(other_mix): # logger.warning(f"Skipping {song_dict} due to empty or invalid audio.") continue if not self.output_mixture: target_clean = target_mix.copy() target_augmented = self.stem_augmentation.apply(target_mix, self.sr) if self.apply_stem_augmentation else target_mix mixture, target_scale, _ = mix_to_target_snr( target_augmented, other_mix, random.uniform(*self.snr_range) ) if not self.output_mixture: target_clean *= target_scale if self.output_mixture: mixture_clean = mixture.copy() mixture_augmented = self.mixture_augmentation.apply(mixture, self.sr) if self.apply_mixture_augmentation else mixture max_val = np.max(np.abs(mixture_augmented)) + 1e-8 mixture_final = mixture_augmented / max_val target_final = (mixture_clean if self.output_mixture else target_clean) / max_val rescale = np.random.uniform(*DEFAULT_GAIN_RANGE) mixture = np.nan_to_num(mixture_final * rescale) target = np.nan_to_num(target_final * rescale) target_length = int(self.clip_duration * self.sr) if target.shape[1] != target_length: target = np.pad(target, (0, target_length - target.shape[1]), mode='constant') else: target = target[:, :target_length] if mixture.shape[1] != target_length: mixture = np.pad(mixture, (0, target_length - mixture.shape[1]), mode='constant') else: mixture = mixture[:, :target_length] return { "mixture": np.nan_to_num(mixture).astype(np.float32), "target": np.nan_to_num(target).astype(np.float32) } # logger.warning(f"No valid audio found for {song_dict}. Skipping.") return self.__getitem__(random.randint(0, len(self.audio_files) - 1)) def __len__(self) -> int: return len(self.audio_files) class InfiniteSampler(Sampler): def __init__(self, dataset: Dataset) -> None: self.dataset_size = len(dataset) self.indexes = list(range(self.dataset_size)) self.reset() def reset(self) -> None: random.shuffle(self.indexes) self.pointer = 0 def __iter__(self): while True: if self.pointer >= self.dataset_size: self.reset() yield self.indexes[self.pointer] self.pointer += 1