mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-15 00:03:06 +00:00
Brandon Thomas
57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
import argparse
|
|
from pathlib import Path
|
|
from concurrent.futures import ProcessPoolExecutor
|
|
from multiprocessing import cpu_count
|
|
|
|
import torchaudio
|
|
from torchaudio.functional import resample
|
|
|
|
from tqdm import tqdm
|
|
|
|
|
|
def process_wav(in_path, out_path, sample_rate):
|
|
wav, sr = torchaudio.load(in_path)
|
|
wav = resample(wav, sr, sample_rate)
|
|
torchaudio.save(out_path, wav, sample_rate)
|
|
return out_path, wav.size(-1) / sample_rate
|
|
|
|
|
|
def preprocess_dataset(args):
|
|
args.out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
futures = []
|
|
executor = ProcessPoolExecutor(max_workers=cpu_count())
|
|
print(f"Resampling audio in {args.in_dir}")
|
|
for in_path in args.in_dir.rglob("*.wav"):
|
|
relative_path = in_path.relative_to(args.in_dir)
|
|
out_path = args.out_dir / relative_path
|
|
out_path.parent.mkdir(parents=True, exist_ok=True)
|
|
futures.append(
|
|
executor.submit(process_wav, in_path, out_path, args.sample_rate)
|
|
)
|
|
|
|
results = [future.result() for future in tqdm(futures)]
|
|
|
|
lengths = {path.stem: length for path, length in results}
|
|
seconds = sum(lengths.values())
|
|
hours = seconds / 3600
|
|
print(f"Wrote {len(lengths)} utterances ({hours:.2f} hours)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Resample an audio dataset.")
|
|
parser.add_argument(
|
|
"in_dir", metavar="in-dir", help="path to the dataset directory.", type=Path
|
|
)
|
|
parser.add_argument(
|
|
"out_dir", metavar="out-dir", help="path to the output directory.", type=Path
|
|
)
|
|
parser.add_argument(
|
|
"--sample-rate",
|
|
help="target sample rate (default 16kHz)",
|
|
type=int,
|
|
default=16000,
|
|
)
|
|
args = parser.parse_args()
|
|
preprocess_dataset(args)
|