| import argparse |
| import glob |
| import os |
| from pathlib import Path |
|
|
| import librosa |
| import torch |
| import torchaudio |
| import multiprocessing as mp |
| import tqdm |
| from tqdm.contrib.concurrent import process_map |
|
|
|
|
| def resample(audio, orig_fs, target_fs=16000): |
| """ |
| Resamples the audio given as input at the target_fs sample rate, if the target sample rate and the |
| original sample rate are different. |
| |
| Args: |
| audio (Tensor): audio to resample |
| orig_fs (int): original sample rate |
| target_fs (int): target sample rate |
| |
| Returns: |
| Tensor: audio resampled |
| """ |
| out = [] |
| for c in range(audio.shape[0]): |
| tmp = audio[c].detach().cpu().numpy() |
| if target_fs != orig_fs: |
| tmp = librosa.resample(tmp, orig_sr=orig_fs, target_sr=target_fs) |
| out.append(torch.from_numpy(tmp)) |
| out = torch.stack(out) |
| return out |
|
|
|
|
| def resample_folder(in_dir, out_dir, target_fs=16000, regex="*.wav"): |
| """ |
| Resamples the audio files contained in the in_dir folder and saves them in out_dir folder |
| |
| Args: |
| in_dir (str): path to audio directory (audio to be resampled) |
| out_dir (str): path to audio resampled directory |
| target_fs (int, optional): target sample rate. Defaults to 16000. |
| regex (str, optional): regular expression for extension of file. Defaults to "*.wav". |
| """ |
| compute = True |
| files = glob.glob(os.path.join(in_dir, regex)) |
| if os.path.exists(out_dir): |
| out_files = glob.glob(os.path.join(out_dir, regex)) |
| if len(files) == len(out_files): |
| compute = False |
| |
| if compute: |
| |
| workers_args = [(f, in_dir, out_dir, target_fs) for f in files] |
| n_workers = min(10, mp.cpu_count()) |
| process_map(_worker_func, workers_args, max_workers=n_workers, chunksize=1) |
| return compute |
|
|
| def _worker_func(input_args): |
| """ |
| Used internally by the pool of multiprocessing workers to resample a given audio file |
| """ |
| f, in_dir, out_dir, target_fs = input_args |
| audio, orig_fs = torchaudio.load(f) |
| audio = resample(audio, orig_fs, target_fs) |
| os.makedirs( |
| Path(os.path.join(out_dir, Path(f).relative_to(Path(in_dir)))).parent, |
| exist_ok=True, |
| ) |
| torchaudio.save( |
| os.path.join(out_dir, Path(f).relative_to(Path(in_dir))), |
| audio, |
| target_fs, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser("Resample a folder recursively") |
| parser.add_argument("--in_dir", type=str) |
| parser.add_argument("--out_dir", type=str) |
| parser.add_argument("--target_fs", default=16000) |
| parser.add_argument("--regex", type=str, default="*.wav") |
|
|
| args = parser.parse_args() |
| resample_folder(args.in_dir, args.out_dir, int(args.target_fs), args.regex) |
|
|
