File size: 2,874 Bytes
3b6a091 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
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 # or thread_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:
# Packing resample_file arguments to the multiprocessing pool
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)
|