o }oi3 @sfddlZddlZddlZddlZddlZddlZddlmZddlm Z m Z m Z m Z ddl Z ddlZddlZddlmZddlmZddlmZmZddlmZddlmZd Zzdd lmZm Z dd lm!Z"Wne#yydd l$m%Z%d ZYnwz ddl&m'Z'd Z(Wn e)e#fyd Z(Ynwd8ddZ*Gddde+Z,Gddde,Z-Gddde,Z.Gddde,Z/Gddde,Z0Gddde,Z1Gddde,Z2Gdd d e,Z3Gd!d"d"e,Z4Gd#d$d$e,Z5Gd%d&d&e,Z6Gd'd(d(e,Z7Gd)d*d*e,Z8e-e.e0e/e1e2e3e4e5e6e7e8d+ Z9d,e:d-e,fd.d/Z;Gd0d1d1e+ZdS):N)NamedTemporaryFile)AnyListOptionalUnion)signal) AudioSegment) collectionsparsers)IterableDataset)loggingT) DictConfig OmegaConf) webdataset)LightningNotInstalledExceptionF) numba_utilsc Cs|r&|dur tdt|\}}}|jdurdn|j}|jdur"dn|j}n t|jdd} | j}| jdur9dn| j}| jdurCdn| j}tj ||||dS)Nz*Expected augmentation dataset but got Noner target_sroffsetduration) TypeErrornextrrrandomsampledata audio_filer from_file) manifestr tarred_audio audio_datasetrfile_idmanifest_entryrr audio_recordr$d/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/parts/preprocessing/perturb.pyread_one_audiosegmentKsr&c@seZdZddZddZdS) PerturbationcCs|SNr$selflengthr$r$r%max_augmentation_length^z$Perturbation.max_augmentation_lengthcCstr()NotImplementedError)r*rr$r$r%perturbar-zPerturbation.perturbN)__name__ __module__ __qualname__r,r/r$r$r$r%r']s r'c@s*eZdZdZd ddZdd Zd d ZdS) SpeedPerturbationa Performs Speed Augmentation by re-sampling the data to a different sampling rate, which does not preserve pitch. Note: This is a very slow operation for online augmentation. If space allows, it is preferable to pre-compute and save the files to augment the dataset. Args: sr: Original sampling rate. resample_type: Type of resampling operation that will be performed. For better speed using `resampy`'s fast resampling method, use `resample_type='kaiser_fast'`. For high-quality resampling, set `resample_type='kaiser_best'`. To use `scipy.signal.resample`, set `resample_type='fft'` or `resample_type='scipy'` min_speed_rate: Minimum sampling rate modifier. max_speed_rate: Maximum sampling rate modifier. num_rates: Number of discrete rates to allow. Can be a positive or negative integer. If a positive integer greater than 0 is provided, the range of speed rates will be discretized into `num_rates` values. If a negative integer or 0 is provided, the full range of speed rates will be sampled uniformly. Note: If a positive integer is provided and the resultant discretized range of rates contains the value '1.0', then those samples with rate=1.0, will not be augmented at all and simply skipped. This is to unnecessary augmentation and increase computation time. Effective augmentation chance in such a case is = `prob * (num_rates - 1 / num_rates) * 100`% chance where `prob` is the global probability of a sample being augmented. rng: Random seed. Default is None ?皙?NcCst||}|dkr td|dvrtd||_||_||_||_|dkr2tj|j|j|jdd|_||_ |r>t |dSddS)N+Minimum sampling rate modifier must be > 0.) kaiser_best kaiser_fastfftscipyzSSupported `resample_type` values are ('kaiser_best', 'kaiser_fast', 'fft', 'scipy')rTendpoint) min ValueError_sr _min_rate _max_rate _num_ratesnplinspace_rates _res_typerseed)r*sr resample_typemin_speed_ratemax_speed_rate num_ratesrngmin_rater$r$r%__init__s zSpeedPerturbation.__init__cC ||jSr(rCr)r$r$r%r, z)SpeedPerturbation.max_augmentation_lengthc Cs|jdkrt|j|j}nt|j}|dkrdSt|j|}zt j j |j |j||j d|_ WdStyT}ztd|jd|d|WYd}~dSd}~ww)Nr?)orig_srrres_typezFailed to resample audio from z to z . Skipping augmentation. Error: )rDruniformrBrCchoicerGintrAlibrosacoreresample_samplesrH Exceptionr warning)r*r speed_ratenew_srer$r$r%r/s  zSpeedPerturbation.perturb)r4r5r6Nr0r1r2__doc__rQr,r/r$r$r$r%r3es   r3c@s*eZdZdZd ddZd d Zd d ZdS)TimeStretchPerturbationa Time-stretch an audio series by a fixed rate while preserving pitch, based on [1]_, [2]_. Note: This is a simplified implementation, intended primarily for reference and pedagogical purposes. It makes no attempt to handle transients, and is likely to produce audible artifacts. References ---------- .. [1] Ellis, D. P. W. "A phase vocoder in Matlab." Columbia University, 2002. ``_ .. [2] librosa.effects.time_stretch ``_ Args: min_speed_rate: Minimum sampling rate modifier. max_speed_rate: Maximum sampling rate modifier. num_rates: Number of discrete rates to allow. Can be a positive or negative integer. If a positive integer greater than 0 is provided, the range of speed rates will be discretized into `num_rates` values. If a negative integer or 0 is provided, the full range of speed rates will be sampled uniformly. Note: If a positive integer is provided and the resultant discretized range of rates contains the value '1.0', then those samples with rate=1.0, will not be augmented at all and simply skipped. This is to avoid unnecessary augmentation and increase computation time. Effective augmentation chance in such a case is = `prob * (num_rates - 1 / num_rates) * 100`% chance where `prob` is the global probability of a sample being augmented. n_fft: Number of fft filters to be computed. rng: Random seed. Default is None r4r5r6NcCsrt||}|dkr td||_||_||_|dkr'tj|j|j|jdd|_|r.t |nFdt ||_ t |d|_ tdtj |j |j d|_tj|j dtjd|_tdtj |j |j d|_tj|j dtjd|_dSt ||_ t |d|_ tdtj |j |j d|_tj|j dtjd|_tdtj |j |j d|_tj|j dtjd|_dS) Nr7r8rTr=rdtype)r?r@rBrCrDrErFrGrrIrZ_n_fft _hop_lengthpi_phi_advance_fastemptyfloat32_scale_buffer_fast_phi_advance_slow_scale_buffer_slow)r*rLrMrNn_fftrOrPr$r$r%rQs,   z TimeStretchPerturbation.__init__cCrRr(rSr)r$r$r%r,rTz/TimeStretchPerturbation.max_augmentation_lengthc Cs|jdkrt|j|j}nt|j}|dkrdS|dkr'd}|j}|j}nd}|j }|j }t |j |}t |j |}tjj|j||d}trRt||||} ntj|||} t tt|j|} tjj| |jj|| d} | |_dS)NrrUrrh)rt hop_length)rjrur+)rDrrXrBrCrYrGrnrqrrrsrZrkrlr[r\stftr^ HAVE_NUMBAr phase_vocoderroundlenistftrj) r*rrafft_multiplier phi_advance scale_bufferrtrurv stft_stretch len_stretch y_stretchr$r$r%r/s.    zTimeStretchPerturbation.perturb)r4r5r6rgNrdr$r$r$r%rfs  ! rfc @sHeZdZdZ      ddededededed ef d d Zd d ZdS)SilencePerturbationa Applies random silence at the start and/or end of the audio. Args: min_start_silence_secs (float): Min start silence level in secs max_start_silence_secs (float): Max start silence level in secs min_end_silence_secs (float): Min end silence level in secs max_end_silence_secs (float): Max end silence level in secs rng (int): Random seed. Default is None value: (float): value representing silence to be added to audio array. rNmin_start_silence_secsmax_start_silence_secsmin_end_silence_secsmax_end_silence_secsrOvaluecCs@||_||_||_||_|rt|nd||_dS||_dSr()_min_start_silence_secs_max_start_silence_secs_min_end_silence_secs_max_end_silence_secsrrI_value)r*rrrrrOrr$r$r%rQ(s   zSilencePerturbation.__init__cCslt|j|j}t|j|j}tt||j f|j }tt||j f|j }t ||j |g|_ dSr() rrXrrrrrEfullrZ sample_rater concatenater^)r*rstart_silence_lenend_silence_lenstartendr$r$r%r/9s zSilencePerturbation.perturb)rrrrNr)r0r1r2refloatrZrQr/r$r$r$r%rs, rc@"eZdZdZd ddZddZdS) GainPerturbationz Applies random gain to the audio. Args: min_gain_dbfs (float): Min gain level in dB max_gain_dbfs (float): Max gain level in dB rng (int): Random seed. Default is None NcC&||_||_|rt|dSddSr()_min_gain_dbfs_max_gain_dbfsrrI)r* min_gain_dbfs max_gain_dbfsrOr$r$r%rQLzGainPerturbation.__init__cCs(t|j|j}|jd|d|_dS)N$@4@)rrXrrr^)r*rgainr$r$r%r/QszGainPerturbation.perturb)rrNr0r1r2rerQr/r$r$r$r%rB rc@s.eZdZdZ      d ddZddZdS) ImpulsePerturbationa Convolves audio with a Room Impulse Response. Args: manifest_path (list): Manifest file for RIRs audio_tar_filepaths (list): Tar files, if RIR audio files are tarred shuffle_n (int): Shuffle parameter for shuffling buffered files from the tar files normalize_impulse (bool): Normalize impulse response to zero mean and amplitude 1 shift_impulse (bool): Shift impulse response to adjust for delay at the beginning rng (int): Random seed. Default is None NFcCs|tj|tgdd|_d|_d|_||_||_d|_ |r-d|_t ||||_t |j|_ ||_ |r:t |j dSddSNTparserindex_by_file_idF)r ASRAudioTextr make_parser _manifest _audiodataset _tarred_audio_normalize_impulse_shift_impulse_data_iteratorAugmentationDatasetiter_rngrrI)r* manifest_pathaudio_tar_filepaths shuffle_nnormalize_impulse shift_impulserOr$r$r%rQcs  zImpulsePerturbation.__init__cCst|j|j|j|jd}|jr!|jt|j}|t t |}n|j}t |j }t t |j dkr9t ddSt|j |d|_ |jrUtt |}|j |d|_ |j d||_ t t |j dkrmt ddS|j t t |j |_ dS)Nrr rz6Zero audio input found, skipping impulse perturbation.rz2Zero audio input found after impulse perturbation.)r&rrrrrsamplesrEmeanmaxabsrzr^r r`r fftconvolverargmax)r*rimpulse impulse_normlen_datamax_indr$r$r%r/{s.   zImpulsePerturbation.perturb)NNrFFNrr$r$r$r%rVs rc@r) ShiftPerturbationa Perturbs audio by shifting the audio in time by a random amount between min_shift_ms and max_shift_ms. The final length of the audio is kept unaltered by padding the audio with zeros. Args: min_shift_ms (float): Minimum time in milliseconds by which audio will be shifted max_shift_ms (float): Maximum time in milliseconds by which audio will be shifted rng (int): Random seed. Default is None @NcCrr() _min_shift_ms _max_shift_msrrI)r* min_shift_ms max_shift_msrOr$r$r%rQrzShiftPerturbation.__init__cCst|j|j}t|d|jkrdSt||jd}|dkr7|jd||j| d<d|jd| <dS|dkrR|j|d|jd| <d|j| d<dSdS)Nir) rrXrrrrrZrr^)r*rshift_ms shift_samplesr$r$r%r/szShiftPerturbation.perturb)rrNrr$r$r$r%rs rc@s\eZdZdZ        ddd Zed d Zd d ZdddZdddZ dddZ dS)NoisePerturbationa Perturbation that adds noise to input audio. Args: manifest_path (str): Manifest file with paths to noise files min_snr_db (float): Minimum SNR of audio after noise is added max_snr_db (float): Maximum SNR of audio after noise is added max_gain_db (float): Maximum gain that can be applied on the noise sample audio_tar_filepaths (list) : Tar files, if noise audio files are tarred shuffle_n (int): Shuffle parameter for shuffling buffered files from the tar files orig_sr (int): Original sampling rate of the noise files rng (int): Random seed. Default is None Nr2r@d>c Cstj|tgdd|_d|_d|_||_d|_|r*d|_t ||||_t |j|_|r1t |nd||_ ||_||_||_dS||_ ||_||_||_dSr)r rr rrrr_orig_srrrrrrIr _min_snr_db _max_snr_db _max_gain_db) r*r min_snr_db max_snr_db max_gain_dbrOrrrVr$r$r%rQs&    zNoisePerturbation.__init__cC|jSr(rr*r$r$r%rVzNoisePerturbation.orig_srcCst|j||j|jdS)Nr)r&rrr)r*rr$r$r%get_one_noise_samplesz&NoisePerturbation.get_one_noise_samplercCs,t|j|j|j|jd}|j|||ddS) Args: data (AudioSegment): audio data ref_mic (int): reference mic index for scaling multi-channel audios r)ref_micN)r&rrrrperturb_with_input_noiser*rrnoiser$r$r%r/szNoisePerturbation.perturbc Cs|j|jkrtd|jd|jdd|kr|jks+ntd|jd|dt|j|j}|rItd|j d |j d |j d |d urP|j }|rltd |j d |j d |j d t d }n|j }|r~|r~tdd S|jdkr|||||}n|||}t||j}td|j |j }|j ||j kr|j|||j d|||jjd|jjdkrtd|jjd|jjd} |j| | |jjd|j7<d S|j|j7_d S)a Args: data (AudioSegment): audio data noise (AudioSegment): noise data data_rms (Union[float, List[float]): rms_db for data input ref_mic (int): reference mic index for scaling multi-channel audios $Found mismatched channels for data ( ) and noise ().r, reference mic ID must be an integer in [0, ), got instead.zEmpty audio segment found for z with offset z and duration .NzEmpty noise segment found for infz>Both data and noise segments are empty. Skipping perturbation.rr7 start_timeend_time) num_channelsr@rrXrris_emptyr r`rrrrms_dbrr?r subsegmentgain_dbr^shaperandint) r*rrdata_rmsrsnr_db noise_rms noise_gain_dbr noise_idxr$r$r%r sH        (z*NoisePerturbation.perturb_with_input_noiserhrcCs|j|jkrtd|jd|jdd|kr|jks+ntd|jd|dt|j|j}|s8|j}|jdkrI|||j||}n||j|}t||j}t d|} t | D]m} td |} td |j } t t | |j} t t t|j | | |j}t|j| |}|d |d 9}|jd|jjdkr|d|jjd}t d|jjd|jd}|j|||jd|7<q`d S) a Args: data (AudioSegment): audio data noise (AudioSegment): noise data data_rms (Union[float, List[float]): rms_db for data input max_noise_dur: (float): max noise duration max_additions (int): number of times for adding noise ref_mic (int): reference mic index for scaling multi-channel audios rrrrrrrrr7rrN)rr@rrXrrrr?rrrangerrZryrrEcopyr^r)r*rrr max_noise_dur max_additionsrrr n_additionsi noise_durr start_sample end_sample noise_samplesrr$r$r%perturb_with_foreground_noiseFs8      "z/NoisePerturbation.perturb_with_foreground_noise)NrrrNNrrrNr)Nrhrr) r0r1r2rerQpropertyrVrr/rrr$r$r$r%rs"    =rc@sxeZdZdZ            dd d Zed dZddZdddZdddZ ddZ d ddZ d!ddZ dS)""NoisePerturbationWithNormalizationa Perturbation that adds noise to input audio, with normalisation to specific decibel level. Also tiles shorter noise samples up to their corresponding clean audio length. Args: manifest_path (str or list): Manifest file with paths to noise files, can be list if using multiple noise sources min_snr_db (float): Minimum SNR of audio after noise is added max_snr_db (float): Maximum SNR of audio after noise is added snr_samples (list): A discrete list of SNRs DBs to sample from when mixing, will be used instead of [min_snr_db,max_snr_db] norm_to_db (float): Will normalise clean, noise, and mixed samples to this DB audio_tar_filepaths (str or list) : Tar files, if noise audio files are tarred, can be list for multiple sources shuffle_n (int): Shuffle parameter for shuffling buffered files from the tar files orig_sr (int): Original sampling rate of the noise files rng (int): Random seed. Default is None shard_strategy (str): if you're using tarred audio and wish to scatter instead of replicate, set this to 'scatter' epsilon (float): minimum value for RMS DB normalisation to avoid divide by zero Nrrrrrr replicate{Gz?c Cs4ddlm}tj|tgdd|_d|_d|_| |_ d|_ |r%t |nd||_ |r|d|_t|tr7|g}t|tr?|g}g}t||D]\}}t|||| | | d}||qF|||r`|nt dd| d|_t|jdkrvtd t|j|_ ||_||_||_t|trt|dkr|nd|_| |_dS) Nr)RandomizedChainDatasetTrF)rank world_sizeshard_strategyi0u)rnd_seedzeNoisePerturbationWithNormalization detected a zero length RandomizedChainDataset, should never happen)'nemo.collections.asr.data.audio_to_textr r rr rrrrrrrrIr isinstancestrziprappendrrz RuntimeErrorrrr _norm_to_dblist _snr_samples_epsilon)r*rrr snr_samples norm_to_dbrOrrrV global_rankrrepsilonr datasetstarred_audio_filepathmanifest_filepathdatasetr$r$r%rQsN       z+NoisePerturbationWithNormalization.__init__cCrr(rrr$r$r%rVrz*NoisePerturbationWithNormalization.orig_srcCs|jrC|jdur tdz t|j\}}}Wnty-t|j|_t|j\}}}Ynw|jdur5dn|j}|jdur?dn|j}n!t |j j dd}|j }|jdurWdn|j}|jduradn|j}tj||||dS)Nz$Expected valid iterator but got Nonerrr)rrrr StopIterationrrrrrrrrrrr)r*rrr!r"rrr#r$r$r%r&s    z8NoisePerturbationWithNormalization.read_one_audiosegmentcCsD||j}|jdkr||j}|jdks |j||||jddS)rr)rrN)r&rrrrrr$r$r%r/s    z*NoisePerturbationWithNormalization.perturb9cCs@|||}|||}d| d}||}||}|||fS)a1 Mixes the clean audio with the noise Args: clean (numpy array): the clean audio data noise (numpy array): the noise audio data snr (float): the SNR value for the mixing norm_to_db (float): the DB value to normalise to before mixing rr)norm_audio_to_db)r*cleanrsnrr noisescalar noisenewlevel noisyspeechr$r$r% snr_mixers  z,NoisePerturbationWithNormalization.snr_mixercCsD|djddd}tt|d|j|}d|d|}||S)z Normalises audio signal to particular db, with some epsilon in-case of divide by zero Args: x (numpy array): input audio signal norm_to_db (float): the db to normalise to rhraxis?rr)rrEwhereiscloser)r*xrrmsscalarr$r$r%r%sz3NoisePerturbationWithNormalization.norm_audio_to_db?cCst|t|kr>|jdkrtt|||jdf}n tt||f}tj||dd}tj||dd}t|t|ks|S)a Tiles the noise array to match the clean audio array, with small silence between the joins Args: clean (numpy array): clean audio data noise (numpy array): noise audio data fs (int): sample rate used by both clean and noise audio data silence_length (float): the amount of silence (in secs) to insert before tiling rrr,)rzndimrEzerosrZrr)r*r&rfssilence_lengthr7 noiseconcatr$r$r%concatenate_noise_samples z;NoisePerturbationWithNormalization.concatenate_noise_samplec Cs$|j|jkrtd|jd|jdd|kr|jks+ntd|jd|d|jr8t|jdd}nt|j|j}|d urUt|j t t j frR|j |n|j }|d ur[|}|j }|j }t|dkrid St|t|kry||||j}|dt|}|j||||d \} } } | |_ d S) a Args: data (AudioSegment): audio data noise (AudioSegment): noise data data_rms (Union[float, List[float]): rms_db for data input ref_mic (int): reference mic index for scaling multi-channel audio, if set to None then each channel will be scaled independently norm_to_db (float): will normalise all audio to this DB rrrrrrrrN)r&rr'r)rr@rrrrXrrrrrrEndarrayr^rzr;rr+) r*rrrrrr data_norm noise_norm_ noisy_snrr$r$r%r"s0 "  z;NoisePerturbationWithNormalization.perturb_with_input_noise) NrrNNNNrrrrr r r)r$)r4)Nrr$) r0r1r2rerQrrVr&r/r+r%r;rr$r$r$r%r vs0 <    r c@r) WhiteNoisePerturbationa5 Perturbation that adds white noise to an audio file in the training dataset. Args: min_level (int): Minimum level in dB at which white noise should be added max_level (int): Maximum level in dB at which white noise should be added rng (int): Random seed. Default is None NcCs0t||_t||_|rtj|dSddSr()rZ min_level max_levelrErrI)r*rDrErOr$r$r%rQYs  zWhiteNoisePerturbation.__init__cCsHtjj|j|jdd}tj|jjdd|d}|j|7_dS)Nint32rirrr)rErrrDrErandnr^r)r*rnoise_level_db noise_signalr$r$r%r/^s zWhiteNoisePerturbation.perturb)rBrCNrr$r$r$r%rAOrrAc@sJeZdZdZ                 dd d ZddZdS)RirAndNoisePerturbationa RIR augmentation with additive foreground and background noise. In this implementation audio data is augmented by first convolving the audio with a Room Impulse Response and then adding foreground noise and background noise at various SNRs. RIR, foreground and background noises should either be supplied with a manifest file or as tarred audio files (faster). Different sets of noise audio files based on the original sampling rate of the noise. This is useful while training a mixed sample rate model. For example, when training a mixed model with 8 kHz and 16 kHz audio with a target sampling rate of 16 kHz, one would want to augment 8 kHz data with 8 kHz noise rather than 16 kHz noise. Args: rir_manifest_path: Manifest file for RIRs rir_tar_filepaths: Tar files, if RIR audio files are tarred rir_prob: Probability of applying a RIR noise_manifest_paths: Foreground noise manifest path min_snr_db: Min SNR for foreground noise max_snr_db: Max SNR for background noise, noise_tar_filepaths: Tar files, if noise files are tarred apply_noise_rir: Whether to convolve foreground noise with a a random RIR orig_sample_rate: Original sampling rate of foreground noise audio max_additions: Max number of times foreground noise is added to an utterance, max_duration: Max duration of foreground noise bg_noise_manifest_paths: Background noise manifest path bg_min_snr_db: Min SNR for background noise bg_max_snr_db: Max SNR for background noise bg_noise_tar_filepaths: Tar files, if noise files are tarred bg_orig_sample_rate: Original sampling rate of background noise audio rng: Random seed. Default is None Nr.rUrrrFr6@rcCs ||_||_||_|rt|ndt|||dd|_d|_d|_|rNi|_t t |D]!}| dur5d}n| |}t ||||||| ||d|j|<q,| |_ | |_ |ri|_t t |D]!}|durhd}n||}t |||||||||d|j|<q_| |_dS)NT)rrrrr)rrrrrV) _rir_prob _noise_prob_bg_noise_probrrIr_rir_perturber_fg_noise_perturbers_bg_noise_perturbersrrzr_max_additions _max_duration_apply_noise_rir)r*rir_manifest_pathrir_probnoise_manifest_paths noise_probrrrir_tar_filepaths rir_shuffle_nnoise_tar_filepathsapply_noise_rirorig_sample_rater max_durationbg_noise_manifest_paths bg_noise_prob bg_min_snr_db bg_max_snr_dbbg_noise_tar_filepathsbg_orig_sample_raterOrrVr$r$r%rQsT z RirAndNoisePerturbation.__init__cCstdd}||jkr|j||j}|jdurQtdd|jkrQ|j}||jvr1t |j }|j|}| |j }|j rE|j||j||||j|jd|jdurtdd|jkr|j}||jvrnt |j }|j|}| |j }|j|||ddSdSdS)Nr7rU)rrr)r)rrXrLrOr/rrPrMrVrkeysrrrTrrSrRrQrNr)r*rprobrrV fg_perturberr bg_perturberr$r$r%r/s.          zRirAndNoisePerturbation.perturb)Nr.NrUrrNrNFNr6rKNrUrrNNNrr$r$r$r%rJds0! ErJc@s"eZdZdZdddZddZdS)TranscodePerturbationak Audio codec augmentation. This implementation uses sox to transcode audio with low rate audio codecs, so users need to make sure that the installed sox version supports the codecs used here (G711 and amr-nb). Args: codecs (List[str]):A list of codecs to be trancoded to. Default is None. rng (int): Random seed. Default is None. NcCsf|rt|nd|dur|ngd|_d|_|dur/|D]}|dvr.td|d|dqdSdS)N)g711amr-nbogg皙?zTranscodePerturbation with z isnot supported. Only z are supported)rrI_codecs att_factorr@)r*codecsrOcodecr$r$r%rQszTranscodePerturbation.__init__c Cstt|j}|dkr|j|}||j}n|j}tdd}t|j| dt dt |j d}|j |dkrgtdd}ttdd }|t dt |d} tjd |jd | d |jd d} nN|j |dkrtdd}ttdd}|t dt |d} tjd |jd | d|jd d} n|j |dkrtdd}tjd |jd|jdd d} tj|jdd} | jd|jjd|_dS)Nrmz.wav)suffixrrrrkz_amr.wavzsox z -V0 -C z2 -t amr-nb - | sox -t amr-nb - -V0 -b 16 -r 16000 T)shellrlz_ogg.wavr5z, -t ogg - | sox -t ogg - -V0 -b 16 -r 16000 rjz _g711.wavz -V0 -r 8000 -c 1 -e a-law z lowpass 3400 highpass 300)r)rErrr^rorsfwritename transposerrrzrnrr subprocess check_outputrrr) r*rrE norm_factor norm_samplesorig_f codec_ind transcoded_fratesrater?new_datar$r$r%r/sB      zTranscodePerturbation.perturbNNrr$r$r$r%ris ric@s*eZdZdZ d ddZdefd d ZdS) RandomSegmentPerturbationa Returns a random segment from input of duration "duration_sec". If duration_sec > input audio length, pad_to_duration determines the outcome. RandomSegmentPerturbation is intended for self-supervised learning. Not for supervised, as extracting corresponding text is not facilitated. Args: duration_sec (float): duration of the segment to be extracted pad_to_duration (bool): zero pad if length of input audio < duration_sec rng: Random seed. Default is None min_rms_db: Minimum RMS db value for the perturbed audio. Default is None max_trials: Maximum number of trials to find a segment with RMS db > min_rms_db. Default is 10 verbose: If True, logs a warning if RMS db < min_rms_db after max_trials. Default is False @@FNrcCsH|dkrtd||_||_||_||_||_|r t|dSddS)Nrzduration_sec should be > 0)r@ _duration_sec_pad_to_duration _min_rms_db _max_trials_verboserrI)r* duration_secpad_to_durationrO min_rms_db max_trialsverboser$r$r%rQ5sz"RandomSegmentPerturbation.__init__rcCs`|j|jkr|js dSd}|j|j|j}|j|dn td|j|j}||j}t |}|j ||d|j dur|j dkrE|j nt|j }d}||j kr||jkrtd|j|j}||j}t |}|j ||d|j dkry|j nt|j }|d7}||j kr||jksV|jr||jkr||j krtd|j d|jd|j ||ddS) Nr7)pad_sizerrrz'Could not find a segment with RMS db > z after z trials.)rrrr num_samplespadrrXrdeepcopyrrrrr?rrr r`)r*rrrrrrtrialr$r$r%r/Bs4      z!RandomSegmentPerturbation.perturb)rFNNrF)r0r1r2rerQrr/r$r$r$r%r#s   r) speed time_stretchrsilencershiftrr> white_noise rir_noise_aug transcode_augrandom_segmentrx perturbationcCs2|tvrtd|dt|d|t|<dS)NzPerturbation with the name z exists. Type of perturbation : r)perturbation_typesreKeyError)rxrr$r$r%register_perturbationqs  rc@s2eZdZd ddZddZddZedd ZdS) AudioAugmentorNcCs.|rt|nd|dur||_dSg|_dSr()rrI _pipeline)r* perturbationsrOr$r$r%rQ{szAudioAugmentor.__init__cCs*|jD]\}}t|kr||qdSr()rrr/)r*segmentrfpr$r$r%r/s   zAudioAugmentor.perturbcCs"|}|jD] \}}||}q|Sr()rr,)r*r+newlenrfrr$r$r%r,s z&AudioAugmentor.max_augmentation_lengthc Csbg}|D]'}|dtvrtd|dqt|d}||d|di|dfq||dS)Naug_typez$%s perturbation not known. Skipping.rfcfgrr$)rr r`r)clsconfigptbsrrr$r$r% from_configs  " zAudioAugmentor.from_configr)r0r1r2rQr/r, classmethodrr$r$r$r%rzs  rrreturnc CsH|durdSt|tr |Sth}trtth}t||vr tdtr.t|tr.tj|dd}t |}g}| D]c\}}| dd}|durOt d|d|d}|dks\|d kr`td z,t|} d t| jvrq||d <d t| jvr}||d <| di|} ||| gWq9t yt d twt|d}|S)a Process list of online data augmentations. Accepts either an AudioAugmentor object with pre-defined augmentations, or a dictionary that points to augmentations that have been defined. If a dictionary is passed, must follow the below structure: Dict[str, Dict[str, Any]]: Which refers to a dictionary of string names for augmentations, defined in `asr/parts/perturb.py`. The inner dictionary may contain key-value arguments of the specific augmentation, along with an essential key `prob`. `prob` declares the probability of the augmentation being applied, and must be a float value in the range [0, 1]. # Example in YAML config file Augmentations are generally applied only during training, so we can add these augmentations to our yaml config file, and modify the behaviour for training and evaluation. ```yaml AudioToSpeechLabelDataLayer: ... # Parameters shared between train and evaluation time train: augmentor: shift: prob: 0.5 min_shift_ms: -5.0 max_shift_ms: 5.0 white_noise: prob: 1.0 min_level: -90 max_level: -46 ... eval: ... ``` Then in the training script, ```python import copy from ruamel.yaml import YAML yaml = YAML(typ="safe") with open(model_config) as f: params = yaml.load(f) # Train Config for Data Loader train_dl_params = copy.deepcopy(params["AudioToTextDataLayer"]) train_dl_params.update(params["AudioToTextDataLayer"]["train"]) del train_dl_params["train"] del train_dl_params["eval"] data_layer_train = nemo_asr.AudioToTextDataLayer( ..., **train_dl_params, ) # Evaluation Config for Data Loader eval_dl_params = copy.deepcopy(params["AudioToTextDataLayer"]) eval_dl_params.update(params["AudioToTextDataLayer"]["eval"]) del eval_dl_params["train"] del eval_dl_params["eval"] data_layer_eval = nemo_asr.AudioToTextDataLayer( ..., **eval_dl_params, ) ``` # Registering your own Augmentations To register custom augmentations to obtain the above convenience of the declaring the augmentations in YAML, you can put additional keys in `perturbation_types` dictionary as follows. ```python from nemo.collections.asr.parts import perturb # Define your own perturbation here class CustomPerturbation(perturb.Perturbation): ... perturb.register_perturbation(name_of_perturbation, CustomPerturbation) ``` Args: augmenter: AudioAugmentor object or dictionary of str -> kwargs (dict) which is parsed and used to initialize an AudioAugmentor. Note: It is crucial that each individual augmentation has a keyword `prob`, that defines a float probability in the the range [0, 1] of this augmentation being applied. If this keyword is not present, then the augmentation is disabled and a warning is logged. Returns: AudioAugmentor object NzCCannot parse augmenter. Must be a dict or an AudioAugmentor object T)resolverfzAugmentation "zW" will not be applied as keyword argument "prob" was not defined for this augmentation.r7rUz-`prob` must be a float value between 0 and 1.rrz,Invalid perturbation name. Allowed values : rr$)rrdictHAVE_OMEGACONG_WEBDATASETr typer@r to_containerrritemsgetrpoprinspect signature parametersrre) augmenterrraugmenter_types augmentations augment_nameaugment_kwargsrfr?augmentation_class augmentationr$r$r%process_augmentationssFP        rc @s`eZdZdZ    ddedeeeefded ed ed ef d d ZddZ ddZ ddZ dS)ra A class that loads tarred audio files and cycles over the files in the dataset. Accepts a single comma-separated JSON manifest file (in the same style as for the AudioToCharDataset/AudioToBPEDataset), as well as the path(s) to the tarball(s) containing the wav files. Each line of the manifest should contain the information for one audio file, including at least the transcript and name of the audio file within the tarball. Valid formats for the audio_tar_filepaths argument include: (1) a single string that can be brace-expanded, e.g. 'path/to/audio.tar' or 'path/to/audio_{1..100}.tar.gz', or (2) a list of file paths that will not be brace-expanded, e.g. ['audio_1.tar', 'audio_2.tar', ...]. Note: For brace expansion in (1), there may be cases where `{x..y}` syntax cannot be used due to shell interference. This occurs most commonly inside SLURM scripts. Therefore we provide a few equivalent replacements. Supported opening braces - { <=> (, [, < and the special tag _OP_. Supported closing braces - } <=> ), ], > and the special tag _CL_. For SLURM based tasks, we suggest the use of the special tags for ease of use. See the WebDataset documentation for more information about accepted data and input formats. rrrr r tar_filepathsrr rrc Csddlm}tj|tgdd|_|||||d}ts t|t t j |dt |t t jddd t d d |j|_dS) Nr)expand_sharded_filepathsTr)rrr)urlsz wav;ogg;flac__key__)audiokeyrr)rrr rr rrrrwds DataPipelineSimpleShardListshuffletarfile_to_samplesrenameto_tuple _loop_offsetsr )r*rrrr rrrr$r$r%rQ)s     zAugmentationDataset.__init__cCs t|jSr()rzrrr$r$r%__len__FrTzAugmentationDataset.__len__csGfddd}||jS)zYThis function is used to iterate through utterances with different offsets for each file.cs(eZdZfddZddZddZdS)zAAugmentationDataset._loop_offsets..TarredAudioLoopOffsetscs"|_||_d|_d|_d|_dSr)iterator collection current_fn current_bytes offset_id)r*rrr$r%rQMs  zJAugmentationDataset._loop_offsets..TarredAudioLoopOffsets.__init__cSs|Sr(r$rr$r$r%__iter__Tr-zJAugmentationDataset._loop_offsets..TarredAudioLoopOffsets.__iter__cSs||jdurt|j\|_|_d|_n$|jj|j}t||jdkr/t|j\|_|_d|_n|jd7_|j|j|jfS)Nrr)rrrrrrmappingrz)r* offset_listr$r$r%__next__Ws zJAugmentationDataset._loop_offsets..TarredAudioLoopOffsets.__next__N)r0r1r2rQrrr$rr$r%TarredAudioLoopOffsetsLs  r)r)r*rrr$rr%rIs z!AugmentationDataset._loop_offsetsc cst|j} z-t|\}}}tjtj|\}}|jj||}|j|}t |} | ||fVWnt yBt|j}Ynwqr() rr rospathsplitextbasenamerrioBytesIOr#) r* audio_iter audio_bytesaudio_filenamerr!r? manifest_idxr"rr$r$r%rgs    zAugmentationDataset.__iter__N)rrrr ) r0r1r2rerrrrZrQrrrr$r$r$r%rs,  r)FN)rr)?rrrrrrztempfilertypingrrrrr[numpyrE soundfilervr<r0nemo.collections.asr.parts.preprocessing.segmentr+nemo.collections.common.parts.preprocessingr r nemo.core.classesr nemo.utilsr r omegaconfr rrrModuleNotFoundErrornemo.utils.exceptionsr nemo.collections.asr.parts.utilsrrw ImportErrorr&objectr'r3rfrrrrrr rArJrirrrrrrrr$r$r$r%s"         Im'P 1Z<?