123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326 |
- import argparse
- import numpy
- import random
- import os
- import time
- import math
- import glob
- import soundfile
- from scipy import signal
- from scipy.io import wavfile
- import soundfile as sf
- import librosa
- import torch
- # import torchaudio
- import torchaudio.transforms as T
- import torch.nn.functional as F
- from tqdm import tqdm
- import multiprocessing as mp
- from multiprocessing import Pool
- from typing import Any, Coroutine, Iterable, List, Tuple
- def loadWAV(filepath, max_audio, evalmode=False, num_eval=10):
- """
- Ignore 'evalmode' and 'num_eval' argument. We not use in this code.
- """
- # Read wav file and convert to torch tensor
- audio, sample_rate = soundfile.read(filepath)
- audiosize = audio.shape[0]
- if audiosize <= max_audio:
- shortage = max_audio - audiosize + 1
- audio = numpy.pad(audio, (0, shortage), 'wrap') # repeat wav to extend the wav length to max length
- audiosize = audio.shape[0]
- if evalmode:
- startframe = numpy.linspace(0,audiosize-max_audio,num=num_eval)
- else:
- startframe = numpy.array([numpy.int64(random.random()*(audiosize-max_audio))])
-
- feats = []
- if evalmode and max_frames == 0:
- feats.append(audio)
- else:
- for asf in startframe:
- feats.append(audio[int(asf):int(asf)+max_audio])
- feat = numpy.stack(feats,axis=0).astype(numpy.float64)
- return feat;
- def loadWav_samplerate(filepath, sample_rate, transform=None):
- waveform, _sample_rate = soundfile.read(filepath)
- if _sample_rate != sample_rate:
- t = torch.from_numpy(waveform).type(torch.float32)
- if transform is None:
- transform = T.Resample(_sample_rate, sample_rate)
- waveform = transform(t).numpy()
- return waveform
- def augment_2ogg(root_path, file_path, sample_rate = 16000, args = None):
-
- sub_path = file_path.rsplit('.', 1)[0].split(root_path)[-1] + '.ogg'
- conv_list = ['ogg']
- conv_list.extend([f'aug{i+1}' for i in range(4)])
- org = None
- for idx, conv in enumerate(conv_list):
- conv_path = root_path + conv + '/' + sub_path
- if os.path.isfile(conv_path):
- continue
-
- os.makedirs(conv_path.rsplit('/', 1)[0], exist_ok=True)
-
- if org is None:
- org = loadWav_samplerate(file_path, sample_rate)
- if conv == 'ogg':
- sf.write(conv_path, org, sample_rate)
- else:
- aug_audio = augmentation.augment_wav_type(org, idx)
- if len(aug_audio.shape) >= 2:
- aug_audio = aug_audio.squeeze(0)
-
- sf.write(conv_path, aug_audio, sample_rate)
-
- class AugmentWAV(object):
- def __init__(self, data_list, dest_dir, musan_path, rir_path, log_interval=100, **kwargs):
- self.data_list = data_list
- self.dest_dir = dest_dir
- self.log_interval = log_interval
- self.noisetypes = ['noise','speech','music']
- self.noisesnr = {'noise':[0,15],'speech':[13,20],'music':[5,15]}
- self.numnoise = {'noise':[1,1], 'speech':[3,7], 'music':[1,1] }
- self.noiselist = {}
- augment_files = glob.glob(os.path.join(musan_path,'*/*/*.wav'))
- for file in augment_files:
- if not file.split('/')[-3] in self.noiselist:
- self.noiselist[file.split('/')[-3]] = []
- self.noiselist[file.split('/')[-3]].append(file)
- self.rir_files = glob.glob(os.path.join(rir_path,'**/*.wav'), recursive=True)
- def augment_wav_type(self, audio, augtype):
- if augtype == 1:
- audio = self.reverberate(audio)
- elif augtype == 2:
- audio = self.additive_noise('music',audio)
- elif augtype == 3:
- audio = self.additive_noise('speech',audio)
- elif augtype == 4:
- audio = self.additive_noise('noise',audio)
- return audio
- async def augment_2ogg(self, file_path, root_path, sample_rate=16000):
- org, samplerate = sf.read(file_path)
- sub_path = file_path.rsplit('.', 1)[0].split(root_path)[-1] + '.ogg'
- ogg_path = root_path + 'ogg/' + sub_path
- if os.path.isfile(ogg_path):
- return
-
- os.makedirs(ogg_path.rsplit('/', 1)[0], exist_ok=True)
- sf.write(ogg_path , org, sample_rate)
- aug_audios = self.augemnt_all(org)
- for idx, aa in enumerate(aug_audios):
- aug_path = f"{root_path}aug{idx+1}/{sub_path}"
- os.makedirs(aug_path.rsplit('/', 1)[0], exist_ok=True)
- if len(aa.shape) >= 2:
- aa = aa.squeeze(0)
- soundfile.write(aug_path, aa, sample_rate)
- def run(self):
- count = 0
- for i, data_path in enumerate(self.data_list):
- count += 1
- filename = os.path.basename(data_path) # get string 'filename.ext' from 'src_dir/filename.ext'
- dest_path = os.path.join(self.dest_dir, filename) # make 'dest_dir/filename.ext'
- audio, sample_rate = soundfile.read(data_path)
- aug_audio = self.augment_wav(audio)
- if len(aug_audio.shape) >= 2:
- aug_audio = aug_audio.squeeze(0)
- soundfile.write(dest_path, aug_audio, sample_rate)
- if i % self.log_interval == 0:
- print(f'{count}/{len(self.data_list)}...')
- def augment_wav(self, audio):
- augtype = random.randint(0,4) # augment audio with 4 noise type randomly.
- if augtype == 1:
- audio = self.reverberate(audio)
- elif augtype == 2:
- audio = self.additive_noise('music',audio)
- elif augtype == 3:
- audio = self.additive_noise('speech',audio)
- elif augtype == 4:
- audio = self.additive_noise('noise',audio)
- return audio
- def augemnt_all(self, audio):
- arr = []
- arr.append(self.reverberate(audio))
- arr.append(self.additive_noise('music',audio))
- arr.append(self.additive_noise('speech',audio))
- arr.append(self.additive_noise('noise',audio))
- return arr
- def additive_noise(self, noisecat, audio):
- max_audio = audio.shape[0]
- clean_db = 10 * numpy.log10(numpy.mean(audio ** 2)+1e-4)
- numnoise = self.numnoise[noisecat]
- noiselist = random.sample(self.noiselist[noisecat], random.randint(numnoise[0],numnoise[1]))
- noises = []
- for noise in noiselist:
- noiseaudio = loadWAV(noise, max_audio, evalmode=False)
- noise_snr = random.uniform(self.noisesnr[noisecat][0],self.noisesnr[noisecat][1])
- noise_db = 10 * numpy.log10(numpy.mean(noiseaudio[0] ** 2)+1e-4)
- noises.append(numpy.sqrt(10 ** ((clean_db - noise_db - noise_snr) / 10)) * noiseaudio)
- return numpy.sum(numpy.concatenate(noises,axis=0),axis=0,keepdims=True) + audio
- def reverberate(self, audio, scale=1.0):
- max_audio = audio.shape[0]
- rir_file = random.choice(self.rir_files)
-
- rir, fs = soundfile.read(rir_file)
- rir = numpy.expand_dims(rir.astype(numpy.float64),0)
- rir = rir / numpy.sqrt(numpy.sum(rir**2))
-
- rirsize = rir.shape[1]
- if rirsize >= max_audio: # If rir size is longer than target audio length
- rir = rir[:,:max_audio]
- if len(rir.shape) >= 3 and rir.shape[2] > 1:
- # rir = rir[:,:,0]
- rir = librosa.to_mono(rir.squeeze().T)
- rir = numpy.expand_dims(rir,0)
- # print(rir.shape, type(rir))
-
- return signal.convolve(numpy.expand_dims(audio,0), rir, mode='full')[:,:max_audio]
- async def augment_func(root_dir, path):
- await augment_2ogg(root_dir, path)
- async def aprogress(tasks: Iterable[Coroutine], **pbar_kws: Any) -> Any:
- if not tasks:
- return -1
-
- pbar = tqdm(asyncio.as_completed(tasks), total=len(tasks), **pbar_kws)
- for task in pbar:
- await task
- pbar.update()
- def str2bool(v):
- if isinstance(v, bool):
- return v
- if v.lower() in ('yes', 'true', 't', 'y', '1'):
- return True
- elif v.lower() in ('no', 'false', 'f', 'n', '0'):
- return False
- else:
- raise argparse.ArgumentTypeError('Boolean value expected.')
- if __name__ == '__main__':
- parser = argparse.ArgumentParser()
- parser.add_argument('--filestxt', default='files.txt', type=str, help='The txt with absolute path of files')
- parser.add_argument('--dest-dir', default='output', type=str, help='The destination save path of augmented audio file')
- parser.add_argument('--musan-path', default='musan_split/', type=str, help='musan file directory')
- parser.add_argument('--rir-path', default='simulated_rirs/', type=str, help='rir file directory')
- parser.add_argument('--log-interval', default=50, type=int)
- parser.add_argument('--root-dir', default='/mnt/data/', type=str, help='setting to the root directory')
- parser.add_argument('--use-list', default=True, type=str2bool, help='use a files.txt')
-
- args = parser.parse_args()
- args.data_list = []
- augmentation = AugmentWAV(**vars(args))
- if args.root_dir[-1] != '/':
- args.root_dir += '/'
- print(args.root_dir)
-
- files_txt = 'files.txt'
- augment_files = []
- if args.use_list and os.path.exists(args.root_dir + files_txt):
- with open(args.root_dir + files_txt, 'r') as f:
- for line in f:
- augment_files.append(line.split('\n')[0])
- else:
- augment_files = glob.glob(os.path.join(args.root_dir,'**/*.wav'), recursive=True)
- if args.use_list:
- with open(args.root_dir + files_txt, 'w') as f:
- for item in augment_files:
- f.write(f'{item}\n')
- print("augment len = ", len(augment_files))
-
- num_cores = mp.cpu_count()
- def work_func(path):
- augment_2ogg(args.root_dir, path, 16000, args)
- try:
- pool = Pool((num_cores)//2)
- with tqdm(total=len(augment_files)) as pbar:
- for _ in tqdm(pool.imap_unordered(work_func, augment_files)):
- pbar.update()
- pool.close()
- pool.join()
- print("Script complete!")
- except KeyboardInterrupt:
- pool.close()
- pool.join()
-
-
|