aug_multi.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326
  1. import argparse
  2. import numpy
  3. import random
  4. import os
  5. import time
  6. import math
  7. import glob
  8. import soundfile
  9. from scipy import signal
  10. from scipy.io import wavfile
  11. import soundfile as sf
  12. import librosa
  13. import torch
  14. # import torchaudio
  15. import torchaudio.transforms as T
  16. import torch.nn.functional as F
  17. from tqdm import tqdm
  18. import multiprocessing as mp
  19. from multiprocessing import Pool
  20. from typing import Any, Coroutine, Iterable, List, Tuple
  21. def loadWAV(filepath, max_audio, evalmode=False, num_eval=10):
  22. """
  23. Ignore 'evalmode' and 'num_eval' argument. We not use in this code.
  24. """
  25. # Read wav file and convert to torch tensor
  26. audio, sample_rate = soundfile.read(filepath)
  27. audiosize = audio.shape[0]
  28. if audiosize <= max_audio:
  29. shortage = max_audio - audiosize + 1
  30. audio = numpy.pad(audio, (0, shortage), 'wrap') # repeat wav to extend the wav length to max length
  31. audiosize = audio.shape[0]
  32. if evalmode:
  33. startframe = numpy.linspace(0,audiosize-max_audio,num=num_eval)
  34. else:
  35. startframe = numpy.array([numpy.int64(random.random()*(audiosize-max_audio))])
  36. feats = []
  37. if evalmode and max_frames == 0:
  38. feats.append(audio)
  39. else:
  40. for asf in startframe:
  41. feats.append(audio[int(asf):int(asf)+max_audio])
  42. feat = numpy.stack(feats,axis=0).astype(numpy.float64)
  43. return feat;
  44. def loadWav_samplerate(filepath, sample_rate, transform=None):
  45. waveform, _sample_rate = soundfile.read(filepath)
  46. if _sample_rate != sample_rate:
  47. t = torch.from_numpy(waveform).type(torch.float32)
  48. if transform is None:
  49. transform = T.Resample(_sample_rate, sample_rate)
  50. waveform = transform(t).numpy()
  51. return waveform
  52. def augment_2ogg(root_path, file_path, sample_rate = 16000, args = None):
  53. sub_path = file_path.rsplit('.', 1)[0].split(root_path)[-1] + '.ogg'
  54. conv_list = ['ogg']
  55. conv_list.extend([f'aug{i+1}' for i in range(4)])
  56. org = None
  57. for idx, conv in enumerate(conv_list):
  58. conv_path = root_path + conv + '/' + sub_path
  59. if os.path.isfile(conv_path):
  60. continue
  61. os.makedirs(conv_path.rsplit('/', 1)[0], exist_ok=True)
  62. if org is None:
  63. org = loadWav_samplerate(file_path, sample_rate)
  64. if conv == 'ogg':
  65. sf.write(conv_path, org, sample_rate)
  66. else:
  67. aug_audio = augmentation.augment_wav_type(org, idx)
  68. if len(aug_audio.shape) >= 2:
  69. aug_audio = aug_audio.squeeze(0)
  70. sf.write(conv_path, aug_audio, sample_rate)
  71. class AugmentWAV(object):
  72. def __init__(self, data_list, dest_dir, musan_path, rir_path, log_interval=100, **kwargs):
  73. self.data_list = data_list
  74. self.dest_dir = dest_dir
  75. self.log_interval = log_interval
  76. self.noisetypes = ['noise','speech','music']
  77. self.noisesnr = {'noise':[0,15],'speech':[13,20],'music':[5,15]}
  78. self.numnoise = {'noise':[1,1], 'speech':[3,7], 'music':[1,1] }
  79. self.noiselist = {}
  80. augment_files = glob.glob(os.path.join(musan_path,'*/*/*.wav'))
  81. for file in augment_files:
  82. if not file.split('/')[-3] in self.noiselist:
  83. self.noiselist[file.split('/')[-3]] = []
  84. self.noiselist[file.split('/')[-3]].append(file)
  85. self.rir_files = glob.glob(os.path.join(rir_path,'**/*.wav'), recursive=True)
  86. def augment_wav_type(self, audio, augtype):
  87. if augtype == 1:
  88. audio = self.reverberate(audio)
  89. elif augtype == 2:
  90. audio = self.additive_noise('music',audio)
  91. elif augtype == 3:
  92. audio = self.additive_noise('speech',audio)
  93. elif augtype == 4:
  94. audio = self.additive_noise('noise',audio)
  95. return audio
  96. async def augment_2ogg(self, file_path, root_path, sample_rate=16000):
  97. org, samplerate = sf.read(file_path)
  98. sub_path = file_path.rsplit('.', 1)[0].split(root_path)[-1] + '.ogg'
  99. ogg_path = root_path + 'ogg/' + sub_path
  100. if os.path.isfile(ogg_path):
  101. return
  102. os.makedirs(ogg_path.rsplit('/', 1)[0], exist_ok=True)
  103. sf.write(ogg_path , org, sample_rate)
  104. aug_audios = self.augemnt_all(org)
  105. for idx, aa in enumerate(aug_audios):
  106. aug_path = f"{root_path}aug{idx+1}/{sub_path}"
  107. os.makedirs(aug_path.rsplit('/', 1)[0], exist_ok=True)
  108. if len(aa.shape) >= 2:
  109. aa = aa.squeeze(0)
  110. soundfile.write(aug_path, aa, sample_rate)
  111. def run(self):
  112. count = 0
  113. for i, data_path in enumerate(self.data_list):
  114. count += 1
  115. filename = os.path.basename(data_path) # get string 'filename.ext' from 'src_dir/filename.ext'
  116. dest_path = os.path.join(self.dest_dir, filename) # make 'dest_dir/filename.ext'
  117. audio, sample_rate = soundfile.read(data_path)
  118. aug_audio = self.augment_wav(audio)
  119. if len(aug_audio.shape) >= 2:
  120. aug_audio = aug_audio.squeeze(0)
  121. soundfile.write(dest_path, aug_audio, sample_rate)
  122. if i % self.log_interval == 0:
  123. print(f'{count}/{len(self.data_list)}...')
  124. def augment_wav(self, audio):
  125. augtype = random.randint(0,4) # augment audio with 4 noise type randomly.
  126. if augtype == 1:
  127. audio = self.reverberate(audio)
  128. elif augtype == 2:
  129. audio = self.additive_noise('music',audio)
  130. elif augtype == 3:
  131. audio = self.additive_noise('speech',audio)
  132. elif augtype == 4:
  133. audio = self.additive_noise('noise',audio)
  134. return audio
  135. def augemnt_all(self, audio):
  136. arr = []
  137. arr.append(self.reverberate(audio))
  138. arr.append(self.additive_noise('music',audio))
  139. arr.append(self.additive_noise('speech',audio))
  140. arr.append(self.additive_noise('noise',audio))
  141. return arr
  142. def additive_noise(self, noisecat, audio):
  143. max_audio = audio.shape[0]
  144. clean_db = 10 * numpy.log10(numpy.mean(audio ** 2)+1e-4)
  145. numnoise = self.numnoise[noisecat]
  146. noiselist = random.sample(self.noiselist[noisecat], random.randint(numnoise[0],numnoise[1]))
  147. noises = []
  148. for noise in noiselist:
  149. noiseaudio = loadWAV(noise, max_audio, evalmode=False)
  150. noise_snr = random.uniform(self.noisesnr[noisecat][0],self.noisesnr[noisecat][1])
  151. noise_db = 10 * numpy.log10(numpy.mean(noiseaudio[0] ** 2)+1e-4)
  152. noises.append(numpy.sqrt(10 ** ((clean_db - noise_db - noise_snr) / 10)) * noiseaudio)
  153. return numpy.sum(numpy.concatenate(noises,axis=0),axis=0,keepdims=True) + audio
  154. def reverberate(self, audio, scale=1.0):
  155. max_audio = audio.shape[0]
  156. rir_file = random.choice(self.rir_files)
  157. rir, fs = soundfile.read(rir_file)
  158. rir = numpy.expand_dims(rir.astype(numpy.float64),0)
  159. rir = rir / numpy.sqrt(numpy.sum(rir**2))
  160. rirsize = rir.shape[1]
  161. if rirsize >= max_audio: # If rir size is longer than target audio length
  162. rir = rir[:,:max_audio]
  163. if len(rir.shape) >= 3 and rir.shape[2] > 1:
  164. # rir = rir[:,:,0]
  165. rir = librosa.to_mono(rir.squeeze().T)
  166. rir = numpy.expand_dims(rir,0)
  167. # print(rir.shape, type(rir))
  168. return signal.convolve(numpy.expand_dims(audio,0), rir, mode='full')[:,:max_audio]
  169. async def augment_func(root_dir, path):
  170. await augment_2ogg(root_dir, path)
  171. async def aprogress(tasks: Iterable[Coroutine], **pbar_kws: Any) -> Any:
  172. if not tasks:
  173. return -1
  174. pbar = tqdm(asyncio.as_completed(tasks), total=len(tasks), **pbar_kws)
  175. for task in pbar:
  176. await task
  177. pbar.update()
  178. def str2bool(v):
  179. if isinstance(v, bool):
  180. return v
  181. if v.lower() in ('yes', 'true', 't', 'y', '1'):
  182. return True
  183. elif v.lower() in ('no', 'false', 'f', 'n', '0'):
  184. return False
  185. else:
  186. raise argparse.ArgumentTypeError('Boolean value expected.')
  187. if __name__ == '__main__':
  188. parser = argparse.ArgumentParser()
  189. parser.add_argument('--filestxt', default='files.txt', type=str, help='The txt with absolute path of files')
  190. parser.add_argument('--dest-dir', default='output', type=str, help='The destination save path of augmented audio file')
  191. parser.add_argument('--musan-path', default='musan_split/', type=str, help='musan file directory')
  192. parser.add_argument('--rir-path', default='simulated_rirs/', type=str, help='rir file directory')
  193. parser.add_argument('--log-interval', default=50, type=int)
  194. parser.add_argument('--root-dir', default='/mnt/data/', type=str, help='setting to the root directory')
  195. parser.add_argument('--use-list', default=True, type=str2bool, help='use a files.txt')
  196. args = parser.parse_args()
  197. args.data_list = []
  198. augmentation = AugmentWAV(**vars(args))
  199. if args.root_dir[-1] != '/':
  200. args.root_dir += '/'
  201. print(args.root_dir)
  202. files_txt = 'files.txt'
  203. augment_files = []
  204. if args.use_list and os.path.exists(args.root_dir + files_txt):
  205. with open(args.root_dir + files_txt, 'r') as f:
  206. for line in f:
  207. augment_files.append(line.split('\n')[0])
  208. else:
  209. augment_files = glob.glob(os.path.join(args.root_dir,'**/*.wav'), recursive=True)
  210. if args.use_list:
  211. with open(args.root_dir + files_txt, 'w') as f:
  212. for item in augment_files:
  213. f.write(f'{item}\n')
  214. print("augment len = ", len(augment_files))
  215. num_cores = mp.cpu_count()
  216. def work_func(path):
  217. augment_2ogg(args.root_dir, path, 16000, args)
  218. try:
  219. pool = Pool((num_cores)//2)
  220. with tqdm(total=len(augment_files)) as pbar:
  221. for _ in tqdm(pool.imap_unordered(work_func, augment_files)):
  222. pbar.update()
  223. pool.close()
  224. pool.join()
  225. print("Script complete!")
  226. except KeyboardInterrupt:
  227. pool.close()
  228. pool.join()