Refactor: add type hints to style_bert_vits2.models.utils

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: