tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Improve: sort wav files before transcription, use pathlib

Browse Source
This commit is contained in:
litagin02
2024-03-01 19:25:21 +09:00
parent 3275124d42
commit f3475287b2
1 changed files with 22 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
transcribe.py
View File

@@ -1,6 +1,7 @@
import argparse import argparse
import os import os
import sys import sys
from pathlib import Path
import yaml import yaml
from faster_whisper import WhisperModel from faster_whisper import WhisperModel
@@ -11,9 +12,9 @@ from common.log import logger
from common.stdout_wrapper import SAFE_STDOUT from common.stdout_wrapper import SAFE_STDOUT
def transcribe(wav_path, initial_prompt=None, language="ja"): def transcribe(wav_path: Path, initial_prompt=None, language="ja"):
segments, _ = model.transcribe( segments, _ = model.transcribe(
wav_path, beam_size=5, language=language, initial_prompt=initial_prompt str(wav_path), beam_size=5, language=language, initial_prompt=initial_prompt
) )
texts = [segment.text for segment in segments] texts = [segment.text for segment in segments]
return "".join(texts) return "".join(texts)
@@ -38,18 +39,18 @@ if __name__ == "__main__":
with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f: with open(os.path.join("configs", "paths.yml"), "r", encoding="utf-8") as f:
path_config: dict[str, str] = yaml.safe_load(f.read()) path_config: dict[str, str] = yaml.safe_load(f.read())
dataset_root = path_config["dataset_root"] dataset_root = Path(path_config["dataset_root"])
model_name = args.model_name model_name = str(args.model_name)
input_dir = os.path.join(dataset_root, model_name, "raw") input_dir = dataset_root / model_name / "raw"
output_file = os.path.join(dataset_root, model_name, "esd.list") output_file = dataset_root / model_name / "esd.list"
initial_prompt = args.initial_prompt initial_prompt = args.initial_prompt
language = args.language language = args.language
device = args.device device = args.device
compute_type = args.compute_type compute_type = args.compute_type
os.makedirs(os.path.dirname(output_file), exist_ok=True) output_file.parent.mkdir(parents=True, exist_ok=True)
logger.info( logger.info(
f"Loading Whisper model ({args.model}) with compute_type={compute_type}" f"Loading Whisper model ({args.model}) with compute_type={compute_type}"
@@ -60,15 +61,18 @@ if __name__ == "__main__":
logger.warning(f"Failed to load model, so use `auto` compute_type: {e}") logger.warning(f"Failed to load model, so use `auto` compute_type: {e}")
model = WhisperModel(args.model, device=device) model = WhisperModel(args.model, device=device)
wav_files = [ # wav_files = [
os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.endswith(".wav") # os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.endswith(".wav")
] # ]
if os.path.exists(output_file): wav_files = [f for f in input_dir.rglob("*.wav") if f.is_file()]
if output_file.exists():
logger.warning(f"{output_file} exists, backing up to {output_file}.bak") logger.warning(f"{output_file} exists, backing up to {output_file}.bak")
if os.path.exists(output_file + ".bak"): # if os.path.exists(output_file + ".bak"):
backup_path = output_file.with_name(output_file.name + ".bak")
if backup_path.exists():
logger.warning(f"{output_file}.bak exists, deleting...") logger.warning(f"{output_file}.bak exists, deleting...")
os.remove(output_file + ".bak") backup_path.unlink()
os.rename(output_file, output_file + ".bak") output_file.rename(backup_path)
if language == "ja": if language == "ja":
language_id = Languages.JP.value language_id = Languages.JP.value
@@ -78,9 +82,11 @@ if __name__ == "__main__":
language_id = Languages.ZH.value language_id = Languages.ZH.value
else: else:
raise ValueError(f"{language} is not supported.") raise ValueError(f"{language} is not supported.")
wav_files = sorted(wav_files, key=lambda x: x.name)
for wav_file in tqdm(wav_files, file=SAFE_STDOUT): for wav_file in tqdm(wav_files, file=SAFE_STDOUT):
file_name = os.path.basename(wav_file)
text = transcribe(wav_file, initial_prompt=initial_prompt, language=language) text = transcribe(wav_file, initial_prompt=initial_prompt, language=language)
with open(output_file, "a", encoding="utf-8") as f: with open(output_file, "a", encoding="utf-8") as f:
f.write(f"{file_name}|{model_name}|{language_id}|{text}\n") f.write(f"{wav_file.name}|{model_name}|{language_id}|{text}\n")
sys.exit(0) sys.exit(0)