From 528d2cc4ba4c2395d83b46583e51f751c44b00ac Mon Sep 17 00:00:00 2001 From: litagin02 Date: Tue, 12 Mar 2024 20:05:32 +0900 Subject: [PATCH] Refactor: typing and pathlib --- slice.py | 48 +++++++++++++++++++++++++++--------------------- style_gen.py | 29 ++++++++++++----------------- transcribe.py | 13 ++++++++----- 3 files changed, 47 insertions(+), 43 deletions(-) diff --git a/slice.py b/slice.py index fb5bb71..99fac19 100644 --- a/slice.py +++ b/slice.py @@ -12,6 +12,8 @@ from style_bert_vits2.logging import logger from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT +# TODO: 並列処理による高速化 + vad_model, utils = torch.hub.load( repo_or_dir="snakers4/silero-vad", model="silero_vad", @@ -23,7 +25,10 @@ vad_model, utils = torch.hub.load( def get_stamps( - audio_file, min_silence_dur_ms: int = 700, min_sec: float = 2, max_sec: float = 12 + audio_file: Path, + min_silence_dur_ms: int = 700, + min_sec: float = 2, + max_sec: float = 12, ): """ min_silence_dur_ms: int (ミリ秒): @@ -42,7 +47,7 @@ def get_stamps( min_ms = int(min_sec * 1000) - wav = read_audio(audio_file, sampling_rate=sampling_rate) + wav = read_audio(str(audio_file), sampling_rate=sampling_rate) speech_timestamps = get_speech_timestamps( wav, vad_model, @@ -56,13 +61,13 @@ def get_stamps( def split_wav( - audio_file, - target_dir="raw", - min_sec=2, - max_sec=12, - min_silence_dur_ms=700, -): - margin = 200 # ミリ秒単位で、音声の前後に余裕を持たせる + audio_file: Path, + target_dir: Path, + min_sec: float = 2, + max_sec: float = 12, + min_silence_dur_ms: int = 700, +) -> tuple[float, int]: + margin: int = 200 # ミリ秒単位で、音声の前後に余裕を持たせる speech_timestamps = get_stamps( audio_file, min_silence_dur_ms=min_silence_dur_ms, @@ -74,10 +79,10 @@ def split_wav( total_ms = len(data) / sr * 1000 - file_name = os.path.basename(audio_file).split(".")[0] - os.makedirs(target_dir, exist_ok=True) + file_name = audio_file.stem + target_dir.mkdir(parents=True, exist_ok=True) - total_time_ms = 0 + total_time_ms: float = 0 count = 0 # タイムスタンプに従って分割し、ファイルに保存 @@ -89,7 +94,7 @@ def split_wav( end_sample = int(end_ms / 1000 * sr) segment = data[start_sample:end_sample] - sf.write(os.path.join(target_dir, f"{file_name}-{i}.wav"), segment, sr) + sf.write(str(target_dir / f"{file_name}-{i}.wav"), segment, sr) total_time_ms += end_ms - start_ms count += 1 @@ -126,20 +131,21 @@ if __name__ == "__main__": ) args = parser.parse_args() - with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f: + with open(Path("configs/paths.yml"), "r", encoding="utf-8") as f: path_config: dict[str, str] = yaml.safe_load(f.read()) dataset_root = path_config["dataset_root"] - input_dir = args.input_dir - output_dir = os.path.join(dataset_root, args.model_name, "raw") - min_sec = args.min_sec - max_sec = args.max_sec - min_silence_dur_ms = args.min_silence_dur_ms + model_name = str(args.model_name) + input_dir = Path(args.input_dir) + output_dir = Path(dataset_root) / model_name / "raw" + min_sec: float = args.min_sec + max_sec: float = args.max_sec + min_silence_dur_ms: int = args.min_silence_dur_ms wav_files = Path(input_dir).glob("**/*.wav") wav_files = list(wav_files) logger.info(f"Found {len(wav_files)} wav files.") - if os.path.exists(output_dir): + if output_dir.exists(): logger.warning(f"Output directory {output_dir} already exists, deleting...") shutil.rmtree(output_dir) @@ -147,7 +153,7 @@ if __name__ == "__main__": total_count = 0 for wav_file in tqdm(wav_files, file=SAFE_STDOUT): time_sec, count = split_wav( - audio_file=str(wav_file), + audio_file=wav_file, target_dir=output_dir, min_sec=min_sec, max_sec=max_sec, diff --git a/style_gen.py b/style_gen.py index f9685cf..b22292d 100644 --- a/style_gen.py +++ b/style_gen.py @@ -1,9 +1,11 @@ import argparse import warnings from concurrent.futures import ThreadPoolExecutor +from typing import Any import numpy as np import torch +from numpy.typing import NDArray from tqdm import tqdm from config import config @@ -11,11 +13,9 @@ from style_bert_vits2.logging import logger from style_bert_vits2.models.hyper_parameters import HyperParameters from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT - warnings.filterwarnings("ignore", category=UserWarning) from pyannote.audio import Inference, Model - model = Model.from_pretrained("pyannote/wespeaker-voxceleb-resnet34-LM") inference = Inference(model, window="whole") device = torch.device(config.style_gen_config.device) @@ -29,11 +29,11 @@ class NaNValueError(ValueError): # 推論時にインポートするために短いが関数を書く -def get_style_vector(wav_path): - return inference(wav_path) +def get_style_vector(wav_path: str) -> NDArray[Any]: + return inference(wav_path) # type: ignore -def save_style_vector(wav_path): +def save_style_vector(wav_path: str): try: style_vec = get_style_vector(wav_path) except Exception as e: @@ -48,20 +48,15 @@ def save_style_vector(wav_path): np.save(f"{wav_path}.npy", style_vec) # `test.wav` -> `test.wav.npy` -def process_line(line): - wavname = line.split("|")[0] +def process_line(line: str): + wav_path = line.split("|")[0] try: - save_style_vector(wavname) + save_style_vector(wav_path) return line, None except NaNValueError: return line, "nan_error" -def save_average_style_vector(style_vectors, filename="style_vectors.npy"): - average_vector = np.mean(style_vectors, axis=0) - np.save(filename, average_vector) - - if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( @@ -71,14 +66,14 @@ if __name__ == "__main__": "--num_processes", type=int, default=config.style_gen_config.num_processes ) args, _ = parser.parse_known_args() - config_path = args.config - num_processes = args.num_processes + config_path: str = args.config + num_processes: int = args.num_processes hps = HyperParameters.load_from_json(config_path) device = config.style_gen_config.device - training_lines = [] + training_lines: list[str] = [] with open(hps.data.training_files, encoding="utf-8") as f: training_lines.extend(f.readlines()) with ThreadPoolExecutor(max_workers=num_processes) as executor: @@ -99,7 +94,7 @@ if __name__ == "__main__": f"Found NaN value in {len(nan_training_lines)} files: {nan_files}, so they will be deleted from training data." ) - val_lines = [] + val_lines: list[str] = [] with open(hps.data.validation_files, encoding="utf-8") as f: val_lines.extend(f.readlines()) diff --git a/transcribe.py b/transcribe.py index 18509c9..f44ecc3 100644 --- a/transcribe.py +++ b/transcribe.py @@ -2,6 +2,7 @@ import argparse import os import sys from pathlib import Path +from typing import Optional import yaml from faster_whisper import WhisperModel @@ -12,7 +13,9 @@ from style_bert_vits2.logging import logger from style_bert_vits2.utils.stdout_wrapper import SAFE_STDOUT -def transcribe(wav_path: Path, initial_prompt=None, language="ja"): +def transcribe( + wav_path: Path, initial_prompt: Optional[str] = None, language: str = "ja" +): segments, _ = model.transcribe( str(wav_path), beam_size=5, language=language, initial_prompt=initial_prompt ) @@ -45,10 +48,10 @@ if __name__ == "__main__": input_dir = dataset_root / model_name / "raw" output_file = dataset_root / model_name / "esd.list" - initial_prompt = args.initial_prompt - language = args.language - device = args.device - compute_type = args.compute_type + initial_prompt: str = args.initial_prompt + language: str = args.language + device: str = args.device + compute_type: str = args.compute_type output_file.parent.mkdir(parents=True, exist_ok=True)