Refactor: typing and pathlib

This commit is contained in:
litagin02
2024-03-12 20:05:32 +09:00
parent f12f20f3da
commit 528d2cc4ba
3 changed files with 47 additions and 43 deletions

View File

@@ -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,

View File

@@ -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())

View File

@@ -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)