From 61e2a1deae543e2403ffa56e97570e11fad4a609 Mon Sep 17 00:00:00 2001 From: tsukumi Date: Sat, 9 Mar 2024 17:26:42 +0000 Subject: [PATCH] Refactor: add type hints to style_bert_vits2.models.utils --- style_bert_vits2/models/utils/__init__.py | 158 +++++++++++++++++----- style_bert_vits2/nlp/__init__.py | 2 +- 2 files changed, 127 insertions(+), 33 deletions(-) diff --git a/style_bert_vits2/models/utils/__init__.py b/style_bert_vits2/models/utils/__init__.py index f488e09..51e19d9 100644 --- a/style_bert_vits2/models/utils/__init__.py +++ b/style_bert_vits2/models/utils/__init__.py @@ -3,28 +3,44 @@ import logging import os import re import subprocess +from pathlib import Path +from typing import Any, Optional, Union import numpy as np import torch +from numpy.typing import NDArray from scipy.io.wavfile import read +from torch.utils.tensorboard import SummaryWriter from style_bert_vits2.logging import logger from style_bert_vits2.models.utils import checkpoints # type: ignore from style_bert_vits2.models.utils import safetensors # type: ignore -MATPLOTLIB_FLAG = False +__is_matplotlib_imported = False def summarize( - writer, - global_step, - scalars={}, - histograms={}, - images={}, - audios={}, - audio_sampling_rate=22050, -): + writer: SummaryWriter, + global_step: int, + scalars: dict[str, float] = {}, + histograms: dict[str, Any] = {}, + images: dict[str, Any] = {}, + audios: dict[str, Any] = {}, + audio_sampling_rate: int = 22050, +) -> None: + """ + 指定されたデータを TensorBoard にまとめて追加する + + Args: + writer (SummaryWriter): TensorBoard への書き込みを行うオブジェクト + global_step (int): グローバルステップ数 + scalars (dict[str, float]): スカラー値の辞書 + histograms (dict[str, Any]): ヒストグラムの辞書 + images (dict[str, Any]): 画像データの辞書 + audios (dict[str, Any]): 音声データの辞書 + audio_sampling_rate (int): 音声データのサンプリングレート + """ for k, v in scalars.items(): writer.add_scalar(k, v, global_step) for k, v in histograms.items(): @@ -35,7 +51,16 @@ def summarize( writer.add_audio(k, v, global_step, audio_sampling_rate) -def is_resuming(dir_path): +def is_resuming(dir_path: Union[str, Path]) -> bool: + """ + 指定されたディレクトリパスに再開可能なモデルが存在するかどうかを返す + + Args: + dir_path: チェックするディレクトリのパス + + Returns: + bool: 再開可能なモデルが存在するかどうか + """ # JP-ExtraバージョンではDURがなくWDがあったり変わるため、Gのみで判断する g_list = glob.glob(os.path.join(dir_path, "G_*.pth")) # d_list = glob.glob(os.path.join(dir_path, "D_*.pth")) @@ -43,13 +68,23 @@ def is_resuming(dir_path): return len(g_list) > 0 -def plot_spectrogram_to_numpy(spectrogram): - global MATPLOTLIB_FLAG - if not MATPLOTLIB_FLAG: +def plot_spectrogram_to_numpy(spectrogram: NDArray[Any]) -> NDArray[Any]: + """ + 指定されたスペクトログラムを画像データに変換する + + Args: + spectrogram (NDArray[Any]): スペクトログラム + + Returns: + NDArray[Any]: 画像データ + """ + + global __is_matplotlib_imported + if not __is_matplotlib_imported: import matplotlib matplotlib.use("Agg") - MATPLOTLIB_FLAG = True + __is_matplotlib_imported = True mpl_logger = logging.getLogger("matplotlib") mpl_logger.setLevel(logging.WARNING) import matplotlib.pylab as plt @@ -63,23 +98,33 @@ def plot_spectrogram_to_numpy(spectrogram): plt.tight_layout() fig.canvas.draw() - data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="") + data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="") # type: ignore data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) plt.close() return data -def plot_alignment_to_numpy(alignment, info=None): - global MATPLOTLIB_FLAG - if not MATPLOTLIB_FLAG: +def plot_alignment_to_numpy(alignment: NDArray[Any], info: Optional[str] = None) -> NDArray[Any]: + """ + 指定されたアライメントを画像データに変換する + + Args: + alignment (NDArray[Any]): アライメント + info (Optional[str]): 画像に追加する情報 + + Returns: + NDArray[Any]: 画像データ + """ + + global __is_matplotlib_imported + if not __is_matplotlib_imported: import matplotlib matplotlib.use("Agg") - MATPLOTLIB_FLAG = True + __is_matplotlib_imported = True mpl_logger = logging.getLogger("matplotlib") mpl_logger.setLevel(logging.WARNING) import matplotlib.pylab as plt - import numpy as np fig, ax = plt.subplots(figsize=(6, 4)) im = ax.imshow( @@ -94,44 +139,93 @@ def plot_alignment_to_numpy(alignment, info=None): plt.tight_layout() fig.canvas.draw() - data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="") + data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="") # type: ignore data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) plt.close() return data -def load_wav_to_torch(full_path): +def load_wav_to_torch(full_path: Union[str, Path]) -> tuple[torch.FloatTensor, int]: + """ + 指定された音声ファイルを読み込み、PyTorch のテンソルに変換して返す + + Args: + full_path (Union[str, Path]): 音声ファイルのパス + + Returns: + tuple[torch.FloatTensor, int]: 音声データのテンソルとサンプリングレート + """ + sampling_rate, data = read(full_path) return torch.FloatTensor(data.astype(np.float32)), sampling_rate -def load_filepaths_and_text(filename, split="|"): +def load_filepaths_and_text(filename: Union[str, Path], split: str = "|") -> list[list[str]]: + """ + 指定されたファイルからファイルパスとテキストを読み込む + + Args: + filename (Union[str, Path]): ファイルのパス + split (str): ファイルの区切り文字 (デフォルト: "|") + + Returns: + list[list[str]]: ファイルパスとテキストのリスト + """ + with open(filename, encoding="utf-8") as f: filepaths_and_text = [line.strip().split(split) for line in f] return filepaths_and_text -def get_logger(model_dir, filename="train.log"): +def get_logger(model_dir_path: Union[str, Path], filename: str = "train.log") -> logging.Logger: + """ + ロガーを取得する + + Args: + model_dir_path (Union[str, Path]): ログを保存するディレクトリのパス + filename (str): ログファイルの名前 (デフォルト: "train.log") + + Returns: + logging.Logger: ロガー + """ + global logger - logger = logging.getLogger(os.path.basename(model_dir)) + logger = logging.getLogger(os.path.basename(model_dir_path)) logger.setLevel(logging.DEBUG) formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s") - if not os.path.exists(model_dir): - os.makedirs(model_dir) - h = logging.FileHandler(os.path.join(model_dir, filename)) + if not os.path.exists(model_dir_path): + os.makedirs(model_dir_path) + h = logging.FileHandler(os.path.join(model_dir_path, filename)) h.setLevel(logging.DEBUG) h.setFormatter(formatter) logger.addHandler(h) return logger -def get_steps(model_path): - matches = re.findall(r"\d+", model_path) +def get_steps(model_path: Union[str, Path]) -> Optional[int]: + """ + モデルのパスからイテレーション番号を取得する + + Args: + model_path (Union[str, Path]): モデルのパス + + Returns: + Optional[int]: イテレーション番号 + """ + + matches = re.findall(r"\d+", model_path) # type: ignore return matches[-1] if matches else None -def check_git_hash(model_dir): +def check_git_hash(model_dir_path: Union[str, Path]) -> None: + """ + モデルのディレクトリに .git ディレクトリが存在する場合、ハッシュ値を比較する + + Args: + model_dir_path (Union[str, Path]): モデルのディレクトリのパス + """ + source_dir = os.path.dirname(os.path.realpath(__file__)) if not os.path.exists(os.path.join(source_dir, ".git")): logger.warning( @@ -143,7 +237,7 @@ def check_git_hash(model_dir): cur_hash = subprocess.getoutput("git rev-parse HEAD") - path = os.path.join(model_dir, "githash") + path = os.path.join(model_dir_path, "githash") if os.path.exists(path): saved_hash = open(path).read() if saved_hash != cur_hash: diff --git a/style_bert_vits2/nlp/__init__.py b/style_bert_vits2/nlp/__init__.py index afc6cb0..683d6d4 100644 --- a/style_bert_vits2/nlp/__init__.py +++ b/style_bert_vits2/nlp/__init__.py @@ -8,7 +8,7 @@ from style_bert_vits2.nlp.symbols import ( ) # __init__.py は配下のモジュールをインポートした時点で実行される -# Pytorch のインポートは重いので、型チェック時以外はインポートしない +# PyTorch のインポートは重いので、型チェック時以外はインポートしない if TYPE_CHECKING: import torch