93 lines
2.5 KiB
Python
93 lines
2.5 KiB
Python
import argparse
|
|
import os
|
|
import shutil
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pandas as pd
|
|
import torch
|
|
from tqdm import tqdm
|
|
|
|
from style_bert_vits2.logging import logger
|
|
|
|
|
|
vad_model, utils = torch.hub.load(
|
|
repo_or_dir="litagin02/silero-vad",
|
|
model="silero_vad",
|
|
onnx=True,
|
|
trust_repo=True,
|
|
)
|
|
|
|
(get_speech_timestamps, _, read_audio, *_) = utils
|
|
|
|
|
|
def get_speech_ratio(audio_file):
|
|
sampling_rate = 16000
|
|
|
|
wav = read_audio(audio_file, sampling_rate=sampling_rate)
|
|
speech_timestamps = get_speech_timestamps(
|
|
wav, vad_model, sampling_rate=sampling_rate
|
|
)
|
|
|
|
speech_dur_ms = 0
|
|
|
|
for ts in speech_timestamps:
|
|
start_ms = ts["start"] / 16
|
|
end_ms = ts["end"] / 16
|
|
speech_dur_ms += end_ms - start_ms
|
|
|
|
total_dur_ms = len(wav) / sampling_rate * 1000
|
|
return speech_dur_ms / total_dur_ms
|
|
|
|
|
|
def process(file: Path):
|
|
speech_ratio = get_speech_ratio(file)
|
|
return file, speech_ratio
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Calculate speech ratio.")
|
|
parser.add_argument(
|
|
"-i", "--input", help="Directory containing audio files", required=True
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if os.path.exists(os.path.join(args.input, "low_speech_ratio")):
|
|
logger.info("Low speech ratio directory already exists, skipping...")
|
|
exit(0)
|
|
|
|
data_dir = Path(args.input)
|
|
wav_files = list(data_dir.glob("*.wav"))
|
|
wav_files.sort()
|
|
|
|
if len(wav_files) < 100:
|
|
logger.warning("Too few files, skipping...")
|
|
exit(0)
|
|
|
|
logger.info(f"Start VAD filtering for {data_dir}...")
|
|
|
|
results = []
|
|
|
|
for wav_file in tqdm(wav_files, file=sys.stdout):
|
|
speech_ratio = get_speech_ratio(wav_file)
|
|
results.append((wav_file, speech_ratio))
|
|
|
|
results_df = pd.DataFrame(results, columns=["file", "speech_ratio"])
|
|
results_df.to_csv(os.path.join(data_dir, "speech_ratio.csv"), index=False)
|
|
|
|
logger.info(f"Speech ratio stats:\n{results_df['speech_ratio'].describe()}")
|
|
threshold = 0.5
|
|
|
|
low_speech_ratio_dir = os.path.join(data_dir, "low_speech_ratio")
|
|
os.makedirs(low_speech_ratio_dir, exist_ok=True)
|
|
|
|
low_speech_files = results_df[results_df["speech_ratio"] < threshold]["file"]
|
|
logger.info(f"Moving {len(low_speech_files)} files to {low_speech_ratio_dir}...")
|
|
for low_speech_file in low_speech_files:
|
|
shutil.move(low_speech_file, low_speech_ratio_dir)
|
|
logger.success("VAD filtering completed.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|