diff --git a/speech_mos.py b/speech_mos.py index 705e657..091148e 100644 --- a/speech_mos.py +++ b/speech_mos.py @@ -1,13 +1,14 @@ import argparse import csv +import re import warnings from pathlib import Path import matplotlib.pyplot as plt +import numpy as np import pandas as pd import torch from tqdm import tqdm -import numpy as np from common.log import logger from common.tts_model import Model @@ -65,6 +66,13 @@ safetensors_files = list(safetensors_files) logger.info(f"There are {len(safetensors_files)} models.") for model_file in tqdm(safetensors_files): + # `test_e10_s1000.safetensors`` -> 1000を取り出す + match = re.search(r"_s(\d+)\.safetensors$", model_file.name) + if match: + step = int(match.group(1)) + else: + logger.warning(f"Step count not found in {model_file.name}, so skip it.") + continue model = get_model(model_file) scores = [] for i, text in enumerate(test_texts): @@ -73,8 +81,6 @@ for model_file in tqdm(safetensors_files): score = predictor(torch.from_numpy(audio).unsqueeze(0), sr).item() scores.append(score) logger.info(f"score: {score}") - # `test_e10_s1000.safetensors`` -> 1000を取り出す - step = int(model_file.stem.split("_")[2][1:]) results.append((model_file.name, step, scores)) del model