wehub-resource-sync
205 lines
7.7 KiB
Python
205 lines
7.7 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import argparse
|
|
import logging
|
|
import os
|
|
import csv
|
|
import tempfile
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
|
|
import torchaudio
|
|
try:
|
|
import webrtcvad
|
|
except ImportError:
|
|
raise ImportError("Please install py-webrtcvad: pip install webrtcvad")
|
|
import pandas as pd
|
|
from tqdm import tqdm
|
|
|
|
from examples.speech_synthesis.preprocessing.denoiser.pretrained import master64
|
|
import examples.speech_synthesis.preprocessing.denoiser.utils as utils
|
|
from examples.speech_synthesis.preprocessing.vad import (
|
|
frame_generator, vad_collector, read_wave, write_wave, FS_MS, THRESHOLD,
|
|
SCALE
|
|
)
|
|
from examples.speech_to_text.data_utils import save_df_to_tsv
|
|
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
PATHS = ["after_denoise", "after_vad"]
|
|
MIN_T = 0.05
|
|
|
|
|
|
def generate_tmp_filename(extension="txt"):
|
|
return tempfile._get_default_tempdir() + "/" + \
|
|
next(tempfile._get_candidate_names()) + "." + extension
|
|
|
|
|
|
def convert_sr(inpath, sr, output_path=None):
|
|
if not output_path:
|
|
output_path = generate_tmp_filename("wav")
|
|
cmd = f"sox {inpath} -r {sr} {output_path}"
|
|
os.system(cmd)
|
|
return output_path
|
|
|
|
|
|
def apply_vad(vad, inpath):
|
|
audio, sample_rate = read_wave(inpath)
|
|
frames = frame_generator(FS_MS, audio, sample_rate)
|
|
frames = list(frames)
|
|
segments = vad_collector(sample_rate, FS_MS, 300, vad, frames)
|
|
merge_segments = list()
|
|
timestamp_start = 0.0
|
|
timestamp_end = 0.0
|
|
# removing start, end, and long sequences of sils
|
|
for i, segment in enumerate(segments):
|
|
merge_segments.append(segment[0])
|
|
if i and timestamp_start:
|
|
sil_duration = segment[1] - timestamp_end
|
|
if sil_duration > THRESHOLD:
|
|
merge_segments.append(int(THRESHOLD / SCALE) * (b'\x00'))
|
|
else:
|
|
merge_segments.append(int((sil_duration / SCALE)) * (b'\x00'))
|
|
timestamp_start = segment[1]
|
|
timestamp_end = segment[2]
|
|
segment = b''.join(merge_segments)
|
|
return segment, sample_rate
|
|
|
|
|
|
def write(wav, filename, sr=16_000):
|
|
# Normalize audio if it prevents clipping
|
|
wav = wav / max(wav.abs().max().item(), 1)
|
|
torchaudio.save(filename, wav.cpu(), sr, encoding="PCM_S",
|
|
bits_per_sample=16)
|
|
|
|
|
|
def process(args):
|
|
# making sure we are requested either denoise or vad
|
|
if not args.denoise and not args.vad:
|
|
log.error("No denoise or vad is requested.")
|
|
return
|
|
|
|
log.info("Creating out directories...")
|
|
if args.denoise:
|
|
out_denoise = Path(args.output_dir).absolute().joinpath(PATHS[0])
|
|
out_denoise.mkdir(parents=True, exist_ok=True)
|
|
if args.vad:
|
|
out_vad = Path(args.output_dir).absolute().joinpath(PATHS[1])
|
|
out_vad.mkdir(parents=True, exist_ok=True)
|
|
|
|
log.info("Loading pre-trained speech enhancement model...")
|
|
model = master64().to(args.device)
|
|
|
|
log.info("Building the VAD model...")
|
|
vad = webrtcvad.Vad(int(args.vad_agg_level))
|
|
|
|
# preparing the output dict
|
|
output_dict = defaultdict(list)
|
|
|
|
log.info(f"Parsing input manifest: {args.audio_manifest}")
|
|
with open(args.audio_manifest, "r") as f:
|
|
manifest_dict = csv.DictReader(f, delimiter="\t")
|
|
for row in tqdm(manifest_dict):
|
|
filename = str(row["audio"])
|
|
|
|
final_output = filename
|
|
keep_sample = True
|
|
n_frames = row["n_frames"]
|
|
snr = -1
|
|
if args.denoise:
|
|
output_path_denoise = out_denoise.joinpath(Path(filename).name)
|
|
# convert to 16khz in case we use a differet sr
|
|
tmp_path = convert_sr(final_output, 16000)
|
|
|
|
# loading audio file and generating the enhanced version
|
|
out, sr = torchaudio.load(tmp_path)
|
|
out = out.to(args.device)
|
|
estimate = model(out)
|
|
estimate = (1 - args.dry_wet) * estimate + args.dry_wet * out
|
|
write(estimate[0], str(output_path_denoise), sr)
|
|
|
|
snr = utils.cal_snr(out, estimate)
|
|
snr = snr.cpu().detach().numpy()[0][0]
|
|
final_output = str(output_path_denoise)
|
|
|
|
if args.vad:
|
|
output_path_vad = out_vad.joinpath(Path(filename).name)
|
|
sr = torchaudio.info(final_output).sample_rate
|
|
if sr in [16000, 32000, 48000]:
|
|
tmp_path = final_output
|
|
elif sr < 16000:
|
|
tmp_path = convert_sr(final_output, 16000)
|
|
elif sr < 32000:
|
|
tmp_path = convert_sr(final_output, 32000)
|
|
else:
|
|
tmp_path = convert_sr(final_output, 48000)
|
|
# apply VAD
|
|
segment, sample_rate = apply_vad(vad, tmp_path)
|
|
if len(segment) < sample_rate * MIN_T:
|
|
keep_sample = False
|
|
print((
|
|
f"WARNING: skip {filename} because it is too short "
|
|
f"after VAD ({len(segment) / sample_rate} < {MIN_T})"
|
|
))
|
|
else:
|
|
if sample_rate != sr:
|
|
tmp_path = generate_tmp_filename("wav")
|
|
write_wave(tmp_path, segment, sample_rate)
|
|
convert_sr(tmp_path, sr,
|
|
output_path=str(output_path_vad))
|
|
else:
|
|
write_wave(str(output_path_vad), segment, sample_rate)
|
|
final_output = str(output_path_vad)
|
|
segment, _ = torchaudio.load(final_output)
|
|
n_frames = segment.size(1)
|
|
|
|
if keep_sample:
|
|
output_dict["id"].append(row["id"])
|
|
output_dict["audio"].append(final_output)
|
|
output_dict["n_frames"].append(n_frames)
|
|
output_dict["tgt_text"].append(row["tgt_text"])
|
|
output_dict["speaker"].append(row["speaker"])
|
|
output_dict["src_text"].append(row["src_text"])
|
|
output_dict["snr"].append(snr)
|
|
|
|
out_tsv_path = Path(args.output_dir) / Path(args.audio_manifest).name
|
|
log.info(f"Saving manifest to {out_tsv_path.as_posix()}")
|
|
save_df_to_tsv(pd.DataFrame.from_dict(output_dict), out_tsv_path)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--audio-manifest", "-i", required=True,
|
|
type=str, help="path to the input manifest.")
|
|
parser.add_argument(
|
|
"--output-dir", "-o", required=True, type=str,
|
|
help="path to the output dir. it will contain files after denoising and"
|
|
" vad"
|
|
)
|
|
parser.add_argument("--vad-agg-level", "-a", type=int, default=2,
|
|
help="the aggresive level of the vad [0-3].")
|
|
parser.add_argument(
|
|
"--dry-wet", "-dw", type=float, default=0.01,
|
|
help="the level of linear interpolation between noisy and enhanced "
|
|
"files."
|
|
)
|
|
parser.add_argument(
|
|
"--device", "-d", type=str, default="cpu",
|
|
help="the device to be used for the speech enhancement model: "
|
|
"cpu | cuda."
|
|
)
|
|
parser.add_argument("--denoise", action="store_true",
|
|
help="apply a denoising")
|
|
parser.add_argument("--vad", action="store_true", help="apply a VAD")
|
|
args = parser.parse_args()
|
|
|
|
process(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|