|
@@ -0,0 +1,326 @@
|
|
|
+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()
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|