# Copyright 2024 The YourMT3 Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Please see the details in the LICENSE file. """augment.py""" import numpy as np import random from collections import defaultdict from typing import Optional, Tuple, Union, Callable, Literal, DefaultDict, Set, Any, Dict, List from utils.note_event_dataclasses import NoteEvent, NoteEventListsBundle from utils.note2event import check_event_len_from_bundle, mix_note_event_lists_bundle, separate_by_subunit_programs_from_note_event_lists_bundle from utils.utils import dict_iterator, extend_dict from copy import deepcopy EPS = 1e-7 DRUM_PROGRAM = 128 UNANNOTATED_PROGRAM = 129 # ------------------------------------------------------------------------------------- # shared augmentation helper functions # ------------------------------------------------------------------------------------- def audio_random_submix_fn(x: np.ndarray, random_amp_range: Optional[List[float]] = None, mask: Optional[np.ndarray] = None, normalize: bool = True, dtype: np.dtype = np.float32) -> Tuple[np.ndarray, np.ndarray]: """ Randomly submix audio. This function supports batch-wise matrix processing. Parameters: - x (np.ndarray): Input audio tensor with shape (b, c, t). - random_amp_range (List[float], optional): A list containing [min_amp, max_amp]. Defaults to [0.6, 1.2]. - mask (np.ndarray, optional): Mask tensor with shape (b, c). Defaults to None. - dtype (np.dtype): Data type for computations. Defaults to np.float32. Returns: - Tuple[np.ndarray, np.ndarray]: Processed audio (stems, mix). """ b, c, t = x.shape if random_amp_range is None: random_amp_range = [0.6, 1.2] if len(random_amp_range) == 2: min_w, max_w = random_amp_range ws = np.random.uniform(min_w, max_w, size=(b, c)).astype(dtype) else: raise ValueError( f"random_amp_range should be a list of two floats, [min_amp, max_amp] or None, but got {random_amp_range}") if mask is not None: ws *= mask # (b, c) processed_audio_stems = x * ws[:, :, np.newaxis] # (b, c, t) processed_audio_mix = np.sum(processed_audio_stems, axis=1, keepdims=True) # (b, 1, t) # Normalize if normalize is True: norm_factors = np.max(np.abs(processed_audio_mix), axis=2, keepdims=True) + EPS # (b, 1, 1) processed_audio_stems /= norm_factors # (b, c, t) processed_audio_mix /= norm_factors # (b, 1, t) else: pass return processed_audio_stems, processed_audio_mix def audio_random_submix_processor(sampled_data: Dict[str, Any], random_amp_range: List[float] = [0.6, 1.2], audio_masks: Optional[List[Optional[np.ndarray]]] = None, update_audio_segments: bool = True, create_processed_audio_array: bool = True) -> None: """Randomly submix audio from sampled data Args: sampled_data: a dictionary containing sampled data. ['audio_segments']: a list of audio segments with length B, each element with shape (1, num_stems, T) random_amp_range: a list of two floats, [min_amp, max_amp] audio_masks: a list of masks. Each mask is binary vector with shape (num_stems,). update_audio_segments: if True (default), update sampled_data["audio_segments"] in-place. create_processed_audio_array: if True (default), create a new key "processed_audio_array" in sampled_data for mix audio. Returns: None (processed audio is stored in sampled_data["processed_audio_array"]) NOTE: - This function creates a new key "processed_audio_array" in sampled_data, in-place of `sampled_data`. - Input audio should exist in sampled_data["audio_segments"]. - The created sampled_data["processed_audio_array"] has shape of (B, 1, T) """ if update_audio_segments is False and create_processed_audio_array is False: raise ValueError("At least one of update_audio_segments and create_processed_audio_mix should be True.") # create a new key "processed_audio" in sampled_data b = len(sampled_data["audio_segments"]) # sub-batch size t = sampled_data["audio_segments"][0].shape[2] # audio length if create_processed_audio_array is True: sampled_data["processed_audio_array"] = np.zeros((b, 1, t), dtype=np.float32) # loop over each audio segment if audio_masks is None: # no audio mask is provided, randomly submix all audio segments for i, audio_segment in enumerate(sampled_data["audio_segments"]): processed_audio_stems, processed_audio_mix = audio_random_submix_fn(x=audio_segment, random_amp_range=random_amp_range, mask=None) if create_processed_audio_array is True: sampled_data["processed_audio_array"][i, :, :] = processed_audio_mix if update_audio_segments is True: sampled_data["audio_segments"][i] = processed_audio_stems else: # audio mask is provided, randomly submix audio segments based on the audio mask for i, (audio_segment, mask) in enumerate(zip(sampled_data["audio_segments"], audio_masks)): processed_audio_stems, processed_audio_mix = audio_random_submix_fn(x=audio_segment, random_amp_range=random_amp_range, mask=mask) if create_processed_audio_array is True: sampled_data["processed_audio_array"][i, :, :] = processed_audio_mix if update_audio_segments is True: sampled_data["audio_segments"][i] = processed_audio_stems def drop_random_stems_from_bundle(sampled_data: Dict[str, Any], prob: float = 0.7) -> None: """ Drop stems with a probability of `prob` from a bundle containing `note_event_segments` and `audio_segments`. It also update `programs`, and add `has_unannotated` info. This function serves as a utility for stem-based data augmentation used by `intra_stem_augment_processor` and `cross_stem_augment_processor`. Args: sampled_data: A dict of sampled data. prob: The probability of dropping stems from the data. Returns: None. The processed data is stored in-place within the `sampled_data` dictionary. Update keys in sampled_data (in-place): sampled_data["note_event_segments"]: NoteEventListsBundle sampled_data["audio_segments"]: NoteEventListsBundle sampled_data["programs_segments"]: a list of list, drum program is 128. updated. sampled_data["has_unannotated_segments"]: a list of bool, True if unannotated program 129 is in use. Newly added. Removed kyes in sampled_data (in-place): all other keys except for the above are removed. Function execution time: 16ms for bsz=36 with single worker """ # Create a deep copy to avoid modifying the original data. note_event_segments = deepcopy(sampled_data["note_event_segments"]) has_unannotated = [] # List of bool, True if unannotated program 129 is in use for i, (has_stems, note_events, tie_note_events, audio_segment, programs, is_drum) in enumerate( zip(sampled_data["has_stems_segments"], note_event_segments['note_events'], note_event_segments['tie_note_events'], sampled_data["audio_segments"], sampled_data["programs_segments"], sampled_data["is_drum_segments"])): # Make sure that programs is np.ndarray if not isinstance(programs, np.ndarray): programs = np.array(programs) if has_stems is True and UNANNOTATED_PROGRAM not in programs: # Get unique and actual presence of instruments. 128 means drums, 129 means unannotated. uniq_programs = np.unique([ne.program if not ne.is_drum else 128 for ne in (tie_note_events + note_events)]) # Debug if DRUM_PROGRAM in uniq_programs: assert DRUM_PROGRAM in programs, "Drum program 128 not in programs" if is_drum.any(): assert DRUM_PROGRAM in programs, "Drum program 128 not in programs" # Vectorized random choice for each unique_program rand_sel_prgs = uniq_programs[np.random.rand(len(uniq_programs)) < prob] if len(rand_sel_prgs) == 0 and len(uniq_programs) != 0: # Make sure at least one program is active rand_sel_prgs = np.random.choice(uniq_programs, size=1) programs_mask = np.isin(programs, rand_sel_prgs).astype(np.int32) drums_mask = programs_mask * is_drum # NOTE: if drums are not annotated as program 128, this would not work properly _programs_in_use = programs[programs_mask == 1] _drum_in_use = np.any(drums_mask == 1) # True if any drum is in use # Drop note_events and tie_note_events in-place note_events[:] = [ ne for ne in note_events if (not ne.is_drum and ne.program in _programs_in_use) or (ne.is_drum and _drum_in_use) ] tie_note_events[:] = [ne for ne in tie_note_events if ne.program in _programs_in_use] # Drop stems from audio_segments, update programs_segments sampled_data["audio_segments"][i] = audio_segment[:, programs_mask == 1, :] sampled_data["programs_segments"][i] = programs[programs_mask == 1] # Create has_unannotated has_unannotated.append(False) elif has_stems is True and UNANNOTATED_PROGRAM in programs: # If unannotated program is included in programs, we only drop 129 with a probability of `prob`. # `note_event_segments` remains the same. # TODO: Actually, we can drop any annoated programs, but current datasets are not the case. uniq_programs = np.unique([ne.program if not ne.is_drum else 128 for ne in (tie_note_events + note_events)]) if np.random.rand() > prob: # keep unannotated program, and this will not allow further cross-stem augmentation. has_unannotated.append(True) else: # drop unannotated program assert UNANNOTATED_PROGRAM not in uniq_programs # 129 is not included here... sampled_data["audio_segments"][i] = audio_segment[:, programs != 129, :] sampled_data["programs_segments"][i] = programs[programs != 129] has_unannotated.append(False) elif has_stems is False and UNANNOTATED_PROGRAM in programs: # No stems, but has unannoted program: cannot be used for cross-stem augmentation. has_unannotated.append(True) else: # No stems, no unannotated program: nothing to do. has_unannotated.append(False) # Update sampled_data in-place sampled_data["note_event_segments"] = note_event_segments sampled_data["has_unannotated_segments"] = has_unannotated # Remove all other keys except for the above, because they are not used in the downstream pipeline. keys_to_remove = ['is_drum_segments', 'has_stems_segments'] for key in keys_to_remove: del sampled_data[key] # ------------------------------------------------------------------------------------- # intra stem augmentation processor # ------------------------------------------------------------------------------------- def intra_stem_augment_processor(sampled_data: Dict[str, Any], random_amp_range: List[float] = [0.6, 1.2], prob: float = 0.7, update_audio_segments: bool = True, submix_audio: bool = True) -> None: """ Intra_stem_augmentation Shape of input: sampled_data: ['note_event_segments']['note_events']: List[List[NoteEvent]] with length B, each element is a list of NoteEvent with length num_notes ['note_event_segments']['tie_note_events']: List[List[NoteEvent]] with length B, each element is a list of NoteEvent with length num_tie_notes ['note_event_segments']['start_times']: List[float] with length B ['audio_segments']: np.ndarray with shape(B, num_stems, T) ['programs_segments']: np.ndarray with shape(num_stems,) ['is_drum_segments']: np.ndarray with shape(num_stems,) ['has_stems_segments']: List[bool] with length B Output (modified in-place): sampled_data: ['note_event_segments']: ['note_events']: ['tie_note_events']: ['start_times']: (not modified) ['audio_segments']: np.ndarray with shape(1, num_stems, T) ['processed_audio_array']: # if submix_audio is True np.ndarray with shape(B, 1, T) ['programs_segments']: List[np.ndarray] with length B, each element is a np.ndarray with shape(num_stems,) ['has_unannotated_segments']: List[bool] with length B Execution time: 27 ms for bsz=36 with single worker, including submix audio """ # Randomly drop stems: # - p (0. < p <= 1.) chances to keep each stem, at least one non-drum is guaranteed to be kept. # - This method modifies the input 'note_event_segments' in-place. drop_random_stems_from_bundle(sampled_data, prob=prob) # Audio processing if submix_audio is True: # Randomly submix audio, and update audio_segments in-place with random amplitude applied. audio_random_submix_processor(sampled_data=sampled_data, random_amp_range=random_amp_range, audio_masks=None, update_audio_segments=True, create_processed_audio_array=True) # mix # assert "processed_audio_array" in sampled_data.keys() else: # NOTE: This is used within the cross-stem augmentation pipeline. pass # ------------------------------------------------------------------------------------- # cross-stem augmentation helper functions # ------------------------------------------------------------------------------------- def combined_survival_and_stop(max_k: int = 5, tau: float = 0.3, alpha: float = 1.0) -> Tuple[np.ndarray, np.ndarray]: """ Compute the survival function and prob_stop for exponential or Weibull distributions based on the value of alpha. - S(k) represents the probability of "surviving" up to k-th trial. - P_stop(k), the stopping probability at trial k is the difference between the survival probabilities at k-1 and k. Parameters: - max_k (int) : Maximum number of trials. k=0, 1, ..., max_k. k=0 means no cross-stem augmentation. - tau (float) : Scale parameter. Represents average time to the first failure for exponential distribution. For Weibull distribution, it influences the spread and shape of the distribution. - alpha (float) : Shape parameter. If alpha=1, the function reduces to exponential distribution. Otherwise, it represents the Weibull distribution. Returns: - survival (array-like) : Computed survival function values. - prob_stop (array-like) : Computed stop probabilities. Example 1: >>> survival_exp, stop_exp = combined_survival_and_stop(max_k=5, tau=0.3, alpha=1.0) Exponential Survival: [1. 0.74081822 0.54881164 0.40656966 0.30119421 0.22313016] Exponential Stop Prob: [0.22313016 0.25918178 0.19200658 0.14224198 0.10537545 0.07806405] Example 2: max_k = 5 survival_exp, stop_exp_03 = combined_survival_and_stop(max_k, 0.3, 1) survival_weibull, stop_weibull = combined_survival_and_stop(max_k, 0.3, 1.5) import matplotlib.pyplot as plt plt.plot(range(max_k+1), list(stop_exp_03), 'o-', label='Exponential (tau=0.3)') plt.plot(range(max_k+1), list(stop_weibull), 's-', label='Weibull (tau=0.3, alpha=1.5)') plt.title("Stop Probabilities"); plt.xlabel("k"); plt.ylabel("Probability") plt.legend(); plt.grid(True); plt.show() References: - Weibull, Waloddi. "A statistical distribution function of wide applicability." Journal of applied mechanics (1951). """ # Generate k values based on max_k k_values = np.arange(max_k + 1) # Calculate survival function if alpha == 1: survival = np.exp(-k_values * tau) else: survival = np.exp(-np.power(k_values * tau, alpha)) # Calculate prob_stop and normalize prob_stop_at_k = -np.diff(np.append(survival, 0.)) return survival, prob_stop_at_k # (max_k+1,), (max_k+1,) def deterministic_random_ux_sampler(prob_stop_at_k, bsz) -> np.ndarray: """ Deterministic random sampler for sampling U\X for cross-stem augmentation. Args: prob_stop_at_k (array-like): Probabilities of stopping at k-th trial. bsz (int) : Batch size. Usually local batch size. Returns: ux_count_per_item (array-like): Number of U\X to sample for each item in the batch. Example: >>> max_k = 5; tau = 0.3; alpha = 1.0; bsz = 20 >>> _, prob_stop_at_k = combined_survival_and_stop(max_k, tau, alpha) prob_stop_at_k: [0.22313016 0.25918178 0.19200658 0.14224198 0.10537545 0.07806405] >>> np.random.choice(np.arange(max_k+1), size=bsz, p=prob_stop_at_k) array([1, 4, 1, 3, 0, 3, 0, 2, 5, 0]) """ ux_count_per_item = np.random.choice(np.arange(len(prob_stop_at_k)), size=bsz, p=prob_stop_at_k) return ux_count_per_item def check_programs_overlap(list_programs: List[np.ndarray], programs: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: """ Check if there is any instrument overlap between two lists of programs. Example: >>> list_programs = np.array([np.array([1,2,3]), np.array([5,6])], dtype=object) >>> print(check_programs_overlap(list_programs, np.array([np.array([1,7])], dtype=object))) # Expected [1] >>> print(check_programs_overlap(list_programs, np.array([np.array([])], dtype=object))) # Expected [] """ list_programs_set = set(item for sublist in list_programs for item in sublist) overlaps = [p for p in programs if p in list_programs_set] uniq_prg_mask = np.array([p not in list_programs_set for p in programs]) return np.array(overlaps), uniq_prg_mask def regroup_program_and_audio_by_minimal_shared_subunits( gathered_programs: List[np.ndarray], gathered_audio_array: List[np.ndarray], max_num_groups: Optional[int] = None ) -> Tuple[List[List[int]], DefaultDict[Tuple[int, ...], List[Tuple[int, int]]]]: # Check if each audio has stems gathered_has_stem = [ audio_array.shape[1] > 1 for programs, audio_array in zip(gathered_programs, gathered_audio_array) ] # Create a dictionary for mapping audio to programs audio2prg = defaultdict(list) for i, programs in enumerate(gathered_programs): for j, value in enumerate(programs): if gathered_has_stem[i] is True: audio2prg[(i, j)].append(value) else: audio2prg[(i, 0)].append(value) grouped_prg2audio = defaultdict(list) for k_tuple, v_list in audio2prg.items(): grouped_prg2audio[tuple(sorted(v_list))].append(k_tuple) # defaultdict(list, # {(61, 69, 71, 72): [(0, 0)], # (128,): [(1, 0)], ...} # Limit the number of groups if max_num_groups is not None: # randomly merge groups while len(grouped_prg2audio) > max_num_groups: # randomly select two groups to merge k1, k2 = random.sample(list(grouped_prg2audio.keys()), 2) grouped_prg2audio[k1].extend(grouped_prg2audio[k2]) del grouped_prg2audio[k2] grouped_programs = list(grouped_prg2audio.keys()) return grouped_programs, grouped_prg2audio # (List[Tuple[int]], DefaultDict[Tuple[int], List[int]]) def audio_random_submix_by_regroup_program_processor(gathered_programs: List[np.ndarray], gathered_audio_array: np.ndarray, submix_random_amp_range: List[float] = [0.9, 1.0], max_num_stems: int = 12) -> Tuple[List[Tuple[int]], np.ndarray]: """Regroup programs into subunit programs, and submix regrouped audio arrays Return: grouped_programs: List[Tuple[int]] submix_audio_array: np.ndarray with shape (1, num_grouped_submix_audio, T) """ # Regroup programs into subunit programs grouped_programs, grouped_prg2audio = regroup_program_and_audio_by_minimal_shared_subunits( gathered_programs, gathered_audio_array, max_num_groups=max_num_stems) # Submix subunit audio arrays, based on the regrouped programs n_frames = gathered_audio_array[0].shape[2] submix_audio_array = np.zeros((1, max_num_stems, n_frames), dtype=np.float32) for i, prgs in enumerate(grouped_programs): audio_ids = grouped_prg2audio[prgs] # id of gathered_audio_array, e.g.:[(i,j),...] if len(audio_ids) == 1: # no need to submix, already subunits src_idx, stem_idx = audio_ids[0] submix_audio_array[:, i, :] = gathered_audio_array[src_idx][:, [stem_idx], :] else: # submix audio from elements of subunit programs _submix_audio_list = [gathered_audio_array[src_idx][:, [stem_idx], :] for (src_idx, stem_idx) in audio_ids] _submix_audio_arr = np.concatenate(_submix_audio_list, axis=1, dtype=np.float32) # (1, C, T) _, _submix_audio_arr = audio_random_submix_fn(_submix_audio_arr, random_amp_range=submix_random_amp_range, normalize=False) submix_audio_array[:, i, :] = _submix_audio_arr return [list(prgs) for prgs in grouped_programs], submix_audio_array # ------------------------------------------------------------------------------------- # cross stem augmentation processor # ------------------------------------------------------------------------------------- def cross_stem_augment_processor( sampled_data: Dict[str, Any], sampled_ids: np.ndarray, get_rand_segments_from_cache_fn: Callable, random_amp_range: List[float] = [0.6, 1.2], stem_iaug_prob: float = 0.7, stem_xaug_policy: Dict = { "max_k": 3, # max number of external sources used for cross-stem augmentations "tau": 0.3, # exponential decay rate for cross-stem augmentation "alpha": 1.0, # shape parameter for Weibull distribution. set 1.0 for exponential. "max_subunit_stems": 12, # the number of subunit stems to be reduced to "p_include_singing": 0.8, # probability of including singing for cross augmented examples. if None, use base probaility. "no_instr_overlap": True, "no_drum_overlap": True, "uhat_intra_stem_augment": True, }, max_l: int = 1024, precomputed_prob_stop_at_k: Optional[np.array] = None, mix_audio: bool = True, create_subunit_note_events: bool = False) -> None: """ Cross-stem augmentation Args: sampled_data: a dictionary containing sampled data. ['note_event_segments']: a list of NoteEventListsBundle with length B ['audio_segments']: a list of audio segments with length B, each element with shape (1, num_stems, T) ['programs_segments']: a list of programs with length B, each element with shape (num_stems,) ['has_unannotated_segments']: a list of bool with length B sampled_ids: a numpy array of sampled ids used in sampled_data. (B,) get_rand_segments_from_cache_fn: a function for getting random segments from cache. random_amp_range: a list of two floats, [min_amp, max_amp] stem_iaug_prob: a float, probability of intra-stem augmentation stem_xaug_policy: a dictionary of cross-stem augmentation policy - max_k (int) : Maximum number of trials. k=0, 1, ..., max_k. k=0 means no cross-stem augmentation. - tau (float) : Scale parameter. Represents average time to the first failure for exponential distribution. For Weibull distribution, it influences the spread and shape of the distribution. - alpha (float) : Shape parameter. If alpha=1, the function reduces to exponential distribution. Otherwise, it represents the Weibull distribution. - max_subunit_stems (int): Maximum number of subunit stems. If larger, they are reduced to this number by submix. Default: 12 - p_include_singing (float): Probability of including singing for cross augmented examples. If None, use base probaility. - no_instr_overlap (bool): If True, do not allow instrument overlap between X and U\X. - no_drum_overlap (bool): If True, do not allow drum overlap between X and U\X. - uhat_intra_stem_augment (bool): If True, apply intra-stem augmentation to U\X. max_l: a int, maximum number of note events in a note event list. Default: 1024 precomputed_prob_stop_at_k: a numpy array of precomputed prob_stop_at_k. If None, it will be computed every time. mix_audio: a bool, if True, mix audio from X and U\X. Default: True create_subunit_note_events: a bool, if True, create subunit note events. This is necessary for multi channel decoder training. Default is False. Returns: None (processed data is stored in-place within the `sampled_data` dictionary) Update keys in sampled_data (in-place): sampled_data["subunit_programs_segments"]: List[List[np.ndarray]], with length B sampled_data["subunit_note_event_segments"]: List[NoteEventListsBundle], with length B sampled_data["subunit_audio_array"]: np.ndarray with shape (B, max_subunit_stems, T) sampled_data["programs_segments"]: List[np.ndarray], with length B sampled_data["note_event_segments"]: NoteEventListsBundle sampled_data["has_unannotated_segments"]: List[bool], with length B sampled_data["processed_audio_array"]: np.ndarray with shape (B, 1, T) Removed kyes in sampled_data (in-place): all other keys except for the above are removed. """ # Setup parameters max_k = stem_xaug_policy["max_k"] tau = stem_xaug_policy["tau"] alpha = stem_xaug_policy.get("alpha", 1.0) max_subunit_stems = stem_xaug_policy.get("max_subunit_stems", 12) p_include_singing = stem_xaug_policy.get("p_include_singing", None) no_instr_overlap = stem_xaug_policy["no_instr_overlap"] no_drum_overlap = stem_xaug_policy["no_drum_overlap"] uhat_intra_stem_augment = stem_xaug_policy["uhat_intra_stem_augment"] bsz = len(sampled_ids) # local batch size n_frames = sampled_data["audio_segments"][0].shape[2] if precomputed_prob_stop_at_k is None: _, prob_stop_at_k = combined_survival_and_stop(max_k, tau, alpha) else: prob_stop_at_k = precomputed_prob_stop_at_k ux_count_per_item = deterministic_random_ux_sampler(prob_stop_at_k, bsz) ux_count_sum = int(np.sum(ux_count_per_item)) # X_in: sampled_data, which we have already applied intra-stem augmentation # U\X: ux_sampled_data, complement of X in U ux_sampled_data, _ = get_rand_segments_from_cache_fn( num_segments=ux_count_sum, use_ordered_read_pos=False, # fully random sampling segments from cache sample_excluding_ids=sampled_ids) # Randomly drop stems from U\X, and update audio stems without submixing audio. if uhat_intra_stem_augment is True: intra_stem_augment_processor(sampled_data=ux_sampled_data, random_amp_range=random_amp_range, prob=stem_iaug_prob, update_audio_segments=True, submix_audio=False) # Loop for creating X_hat iter_ux = iter( zip( ux_sampled_data['audio_segments'], dict_iterator(ux_sampled_data['note_event_segments']), ux_sampled_data['programs_segments'], ux_sampled_data['has_unannotated_segments'], )) iter_x_in = iter( zip( sampled_data['audio_segments'], dict_iterator(sampled_data['note_event_segments']), sampled_data['programs_segments'], sampled_data['has_unannotated_segments'], )) x_hat = { "subunit_programs_segments": [], # List[List[np.ndarray]], with length B "subunit_note_event_segments": [], # List[NoteEventListsBundle], with length B "subunit_audio_array": np.zeros((bsz, max_subunit_stems, n_frames), dtype=np.float32), # (B, max_submix_stems, T) "programs_segments": [], # List[np.ndarray], with length B "note_event_segments": { "note_events": [], "tie_note_events": [], "start_times": [] }, # NoteEventListsBundle "has_unannotated_segments": [], # List[bool], with length B "processed_audio_array": np.zeros((bsz, 1, n_frames), dtype=np.float32), # mixed audio array, B, 1, T) } for i, (audio_array, ne_bundle, programs, has_unannotated) in enumerate(iter_x_in): num_ux_samples = ux_count_per_item[i] if num_ux_samples > 0 and has_unannotated is False: # gather the main source and k external sources gathered_programs = [programs] gathered_ne_bundle = ne_bundle # mutable, but ok because `dict_iterator` yields new dict gathered_audio_array = [audio_array] for k in range(num_ux_samples): # Get next external source ex_audio_array, ex_ne_bundle, ex_programs, ex_has_unannotated = next(iter_ux) ex_prg_mask = None # None: no need to mask external programs ex_has_stem = bool(ex_audio_array.shape[1] > 1) """Criteria for skipping sources""" if ex_has_unannotated is True: continue """Criteria for instrument overlap and drum overlap """ instr_overlap, uniq_ex_prg_mask = check_programs_overlap(gathered_programs, ex_programs) if no_instr_overlap is True and len(instr_overlap) > 0: if np.any(uniq_ex_prg_mask) and ex_has_stem is True: # mask out non-unique external programs ex_prg_mask = uniq_ex_prg_mask else: # print(i, k, num_ux_samples, ex_programs, # 'Warning: no unique external programs, skip this source') continue # no unique external programs, skip this source else: # programs is already unique or don't care about overlap pass if no_drum_overlap is True and no_instr_overlap is False and DRUM_PROGRAM in instr_overlap: non_drum_ex_prg_mask = np.array([prg != DRUM_PROGRAM for prg in ex_programs]) if np.any(non_drum_ex_prg_mask): # mask only drum external programs ex_prg_mask = non_drum_ex_prg_mask else: # print(i, k, num_ux_samples, ex_programs, # 'Warning: no non-drum external programs, skip this source') continue # drum overlapped, but no non-drum programs, skip this source else: pass """Criteria for stopping iteration with respect to max length""" if check_event_len_from_bundle(gathered_ne_bundle, ex_ne_bundle, max_len=max_l) is False: # print(i, k, num_ux_samples, 'Warning: max length reached, stop iteration') break # Apply mask and gather if ex_prg_mask is None: gathered_programs.append(ex_programs) extend_dict(gathered_ne_bundle, ex_ne_bundle) gathered_audio_array.append(ex_audio_array) else: # apply mask to external programs, and add to list ex_programs = ex_programs[ex_prg_mask] gathered_programs.append(ex_programs) # drop note_events with masked programs, and extend dictionary _ex_has_drum = np.any(ex_programs == DRUM_PROGRAM) ex_ne_bundle["note_events"][0] = [ ne for ne in ex_ne_bundle["note_events"][0] if (not ne.is_drum and ne.program in ex_programs) or (ne.is_drum and _ex_has_drum) ] ex_ne_bundle["tie_note_events"][0] = [ ne for ne in ex_ne_bundle["tie_note_events"][0] if ne.program in ex_programs ] extend_dict(gathered_ne_bundle, ex_ne_bundle) # apply mask to external audio_array, and add to list gathered_audio_array.append(ex_audio_array[:, ex_prg_mask, :]) # print(gathered_programs) # Regroup gathered programs, and cresate submix by subunits programs subunit_programs, subunit_audio_array = audio_random_submix_by_regroup_program_processor( gathered_programs, gathered_audio_array, max_num_stems=max_subunit_stems) mixed_ne_bundle = mix_note_event_lists_bundle(gathered_ne_bundle, sort=True, start_time_to_zero=True, use_deepcopy=True) #False) if create_subunit_note_events is True: subunit_ne_bundle = separate_by_subunit_programs_from_note_event_lists_bundle(mixed_ne_bundle, subunit_programs, start_time_to_zero=False, sort=True) else: subunit_ne_bundle = None x_hat["subunit_note_event_segments"].append(subunit_ne_bundle) x_hat["subunit_programs_segments"].append(subunit_programs) x_hat["subunit_audio_array"][i, :subunit_audio_array.shape[1], :] = subunit_audio_array # (B, C, T) x_hat["programs_segments"].append(np.concatenate(gathered_programs, axis=0)) extend_dict(x_hat["note_event_segments"], mixed_ne_bundle) x_hat["has_unannotated_segments"].append(has_unannotated) else: num_stems = audio_array.shape[1] if num_stems > max_subunit_stems: # If num_stems exceeds max_subunit_stems, randomly select max_subunit_stems stems subunit_programs, subunit_audio_array = audio_random_submix_by_regroup_program_processor( [programs], [audio_array], max_num_stems=max_subunit_stems) else: subunit_programs = [programs] subunit_audio_array = audio_array x_hat["subunit_programs_segments"].append(subunit_programs) x_hat["subunit_audio_array"][i, :subunit_audio_array.shape[1], :] = subunit_audio_array if create_subunit_note_events is True: subunit_ne_bundle = separate_by_subunit_programs_from_note_event_lists_bundle(ne_bundle, subunit_programs, start_time_to_zero=True, sort=True) else: subunit_ne_bundle = None x_hat["subunit_note_event_segments"].append(subunit_ne_bundle) x_hat["programs_segments"].append(programs) extend_dict(x_hat["note_event_segments"], ne_bundle) x_hat["has_unannotated_segments"].append(has_unannotated) # Mix subunit audio and update subunit audio arrays if mix_audio is True: amp_applied_stem_arr, mix_audio_arr = audio_random_submix_fn(x_hat["subunit_audio_array"], random_amp_range=random_amp_range, mask=None, normalize=True) x_hat["subunit_audio_array"] = amp_applied_stem_arr # (B, C, T) x_hat["processed_audio_array"] = mix_audio_arr # (B, 1, T) # Update sampled_data in-place sampled_data["subunit_programs_segments"] = x_hat["subunit_programs_segments"] sampled_data["subunit_note_event_segments"] = x_hat["subunit_note_event_segments"] sampled_data["subunit_audio_array"] = x_hat["subunit_audio_array"] sampled_data["programs_segments"] = x_hat["programs_segments"] sampled_data["note_event_segments"] = x_hat["note_event_segments"] sampled_data["has_unannotated_segments"] = x_hat["has_unannotated_segments"] sampled_data["processed_audio_array"] = x_hat["processed_audio_array"] del sampled_data["audio_segments"]