From cca9a06c8cc57b721bac479ed4d23eeeb4770f42 Mon Sep 17 00:00:00 2001 From: BenAAndrew Date: Thu, 25 Nov 2021 21:43:59 +0000 Subject: [PATCH] Enable other Silero languages --- alphabets/English.txt | 33 +++++++++++++++++++++++++++++++++ alphabets/German.txt | 36 ++++++++++++++++++++++++++++++++++++ alphabets/Spanish.txt | 34 ++++++++++++++++++++++++++++++++++ application/views.py | 36 ++++++++++++++++++++++++------------ dataset/create_dataset.py | 4 ++-- dataset/transcribe.py | 17 ++++++----------- tests/test_dataset.py | 12 +++++------- tests/test_synthesis.py | 4 ++-- 8 files changed, 142 insertions(+), 34 deletions(-) create mode 100644 alphabets/English.txt create mode 100644 alphabets/German.txt create mode 100644 alphabets/Spanish.txt diff --git a/alphabets/English.txt b/alphabets/English.txt new file mode 100644 index 0000000..46aa35e --- /dev/null +++ b/alphabets/English.txt @@ -0,0 +1,33 @@ +# Each line in this file represents the Unicode codepoint (UTF-8 encoded) +# associated with a numeric label. +# A line that starts with # is a comment. You can escape it with \# if you wish +# to use '#' as a label. + +a +b +c +d +e +f +g +h +i +j +k +l +m +n +o +p +q +r +s +t +u +v +w +x +y +z +' +# The last (non-comment) line needs to end with a newline. diff --git a/alphabets/German.txt b/alphabets/German.txt new file mode 100644 index 0000000..919eb91 --- /dev/null +++ b/alphabets/German.txt @@ -0,0 +1,36 @@ +# Each line in this file represents the Unicode codepoint (UTF-8 encoded) +# associated with a numeric label. +# A line that starts with # is a comment. You can escape it with \# if you wish +# to use '#' as a label. + +a +b +c +d +e +f +g +h +i +j +k +l +m +n +o +p +q +r +s +t +u +v +w +x +y +z +ä +ö +ü +' +# The last (non-comment) line needs to end with a newline. diff --git a/alphabets/Spanish.txt b/alphabets/Spanish.txt new file mode 100644 index 0000000..9fe5eac --- /dev/null +++ b/alphabets/Spanish.txt @@ -0,0 +1,34 @@ +# Each line in this file represents the Unicode codepoint (UTF-8 encoded) +# associated with a numeric label. +# A line that starts with # is a comment. You can escape it with \# if you wish +# to use '#' as a label. + +a +b +c +d +e +f +g +h +i +j +k +l +m +n +o +p +q +r +s +t +u +v +w +x +y +z +' +ñ +# The last (non-comment) line needs to end with a newline. diff --git a/application/views.py b/application/views.py index b83d88e..1297cf4 100644 --- a/application/views.py +++ b/application/views.py @@ -27,8 +27,7 @@ from dataset.create_dataset import ( from dataset.clip_generator import CHARACTER_ENCODING, add_suffix from dataset.extend_existing_dataset import extend_existing_dataset from dataset.analysis import get_total_audio_duration, validate_dataset -from dataset.transcribe import create_transcription_model -from training import DEFAULT_ALPHABET +from dataset.transcribe import Silero, DeepSpeech, SILERO_LANGUAGES from training.train import train, TRAINING_PATH from training.utils import get_available_memory, get_batch_size, load_symbols, generate_timelapse_gif from synthesis.synthesize import load_model, synthesize @@ -38,6 +37,7 @@ from flask import redirect, render_template, request, send_file URLS = {"/": "Build dataset", "/train": "Train", "/synthesis-setup": "Synthesis"} +ALPHABET_FOLDER = "alphabets" TEXT_FILE = "text.txt" SUBTITLE_FILE = "sub.srt" CHECKPOINT_FOLDER = "checkpoints" @@ -46,7 +46,6 @@ RESULTS_FILE = "out.wav" TEMP_DATASET_UPLOAD = "temp.zip" TRANSCRIPTION_MODEL = "model.pbmm" ALPHABET_FILE = "alphabet.txt" -ENGLISH_LANGUAGE = "English" model = None vocoder = None @@ -54,8 +53,12 @@ symbols = None def get_languages(): - custom_models = {language: os.path.isfile(os.path.join(paths["languages"], language, TRANSCRIPTION_MODEL)) for language in os.listdir(paths["languages"])} - return {**{ENGLISH_LANGUAGE: True}, **custom_models} + silero_languages = {language: True for language in SILERO_LANGUAGES} + custom_models = { + language: os.path.isfile(os.path.join(paths["languages"], language, TRANSCRIPTION_MODEL)) + for language in os.listdir(paths["languages"]) + } + return {**silero_languages, **custom_models} def get_checkpoints(): @@ -95,10 +98,11 @@ def create_dataset_post(): combine_clips = request.form.get("combine_clips") is not None min_length = float(request.form["min_length"]) max_length = float(request.form["max_length"]) - transcription_model_path = ( - os.path.join(paths["languages"], language, TRANSCRIPTION_MODEL) if language != ENGLISH_LANGUAGE else None + transcription_model = ( + Silero(language) + if language in SILERO_LANGUAGES + else DeepSpeech(os.path.join(paths["languages"], language, TRANSCRIPTION_MODEL)) ) - transcription_model = create_transcription_model(transcription_model_path) text_file = SUBTITLE_FILE if request.files["text_file"].filename.endswith(".srt") else TEXT_FILE if request.form["name"]: @@ -198,7 +202,11 @@ def get_train(): @app.route("/train", methods=["POST"]) def train_post(): language = request.form["language"] - alphabet_path = os.path.join(paths["languages"], language, ALPHABET_FILE) if language != ENGLISH_LANGUAGE else None + alphabet_path = ( + os.path.join(ALPHABET_FOLDER, f"{language}.txt") + if language in SILERO_LANGUAGES + else os.path.join(paths["languages"], language, ALPHABET_FILE) + ) dataset_name = request.form["dataset"] epochs = request.form["epochs"] batch_size = request.form["batch_size"] @@ -277,8 +285,12 @@ def synthesis_setup_post(): vocoder = Hifigan(model_path, model_config_path) dataset_name = request.form["model"] language = request.form["language"] - alphabet_path = os.path.join(paths["languages"], language, ALPHABET_FILE) - symbols = load_symbols(alphabet_path) if language != ENGLISH_LANGUAGE else DEFAULT_ALPHABET + alphabet_path = ( + os.path.join(ALPHABET_FOLDER, f"{language}.txt") + if language in SILERO_LANGUAGES + else os.path.join(paths["languages"], language, ALPHABET_FILE) + ) + symbols = load_symbols(alphabet_path) checkpoint_folder = os.path.join(paths["models"], dataset_name) checkpoint = os.path.join(checkpoint_folder, request.form["checkpoint"]) model = load_model(checkpoint) @@ -476,7 +488,7 @@ def upload_language(): language = request.values["name"] language_dir = os.path.join(paths["languages"], language) os.makedirs(language_dir, exist_ok=True) - if(request.files["model"]): + if request.files["model"]: request.files["model"].save(os.path.join(language_dir, TRANSCRIPTION_MODEL)) request.files["alphabet"].save(os.path.join(language_dir, ALPHABET_FILE)) return redirect("/settings") diff --git a/dataset/create_dataset.py b/dataset/create_dataset.py index 0ed7a3e..cff6d79 100644 --- a/dataset/create_dataset.py +++ b/dataset/create_dataset.py @@ -10,7 +10,7 @@ sys.path.append(dirname(dirname(abspath(__file__)))) from dataset.audio_processing import convert_audio from dataset.clip_generator import clip_generator, MIN_LENGTH, MAX_LENGTH from dataset.analysis import save_dataset_info -from dataset.transcribe import create_transcription_model +from dataset.transcribe import Silero AUDIO_FOLDER = "wavs" @@ -102,4 +102,4 @@ if __name__ == "__main__": parser.add_argument("-i", "--info_path", help="Path to save info file", type=str, default="info.json") args = parser.parse_args() - create_dataset(**vars(args), transcription_model=create_transcription_model()) + create_dataset(**vars(args), transcription_model=Silero()) diff --git a/dataset/transcribe.py b/dataset/transcribe.py index bc963ed..156d71d 100644 --- a/dataset/transcribe.py +++ b/dataset/transcribe.py @@ -14,6 +14,9 @@ import omegaconf # noqa from dataset.silero_utils import init_jit_model +SILERO_LANGUAGES = {"English": "en", "German": "de", "Spanish": "es"} + + class TranscriptionModel(ABC): @abstractmethod def load_audio(self, path): @@ -82,7 +85,7 @@ class Silero(TranscriptionModel): Credit: https://github.com/snakers4/silero-models """ - def __init__(self, language="en"): + def __init__(self, language="English"): self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = os.path.join(getattr(sys, "_MEIPASS", ""), "en_v5.jit") if os.path.isfile(model): @@ -91,7 +94,7 @@ class Silero(TranscriptionModel): self.model, self.decoder, _ = torch.hub.load( repo_or_dir="snakers4/silero-models", model="silero_stt", - language=language, + language=SILERO_LANGUAGES[language], device=self.device, ) @@ -114,20 +117,12 @@ class Silero(TranscriptionModel): return self.decoder(example.cpu()) -def create_transcription_model(model_path=None): - if model_path: - return DeepSpeech(model_path) - # If no model path, default to English Sliero - else: - return Silero() - - if __name__ == "__main__": """Transcribe a clip""" parser = argparse.ArgumentParser(description="Transcribe a clip") parser.add_argument("-i", "--input_path", help="Path to audio file", type=str, required=True) args = parser.parse_args() - model = create_transcription_model() + model = Silero() text = model.transcribe(args.input_path) print("Text: ", text) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 1952b31..40541f6 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -12,7 +12,7 @@ from dataset.clip_generator import generate_clips_from_subtitles, clip_combiner from dataset.create_dataset import create_dataset from dataset.extend_existing_dataset import extend_existing_dataset from dataset.utils import similarity, add_suffix -from dataset.transcribe import create_transcription_model, TranscriptionModel, DeepSpeech, Silero +from dataset.transcribe import TranscriptionModel, DeepSpeech, Silero TEXT = "the examination and testimony of the experts enabled the commission to conclude that five shots may have been fired" @@ -328,18 +328,16 @@ def test_validate_dataset(): # Transcription def test_deepspeech(): model_path = os.path.join("test_samples", "english.pbmm") - deepspeech = create_transcription_model(model_path) - assert isinstance(deepspeech, DeepSpeech) + transcription_model = DeepSpeech(model_path) audio_path = os.path.join("test_samples", "audio.wav") - transcription = deepspeech.transcribe(audio_path) + transcription = transcription_model.transcribe(audio_path) assert similarity(TEXT, transcription) > MIN_SYNTHESIS_SCORE def test_silero(): - silero = create_transcription_model() - assert isinstance(silero, Silero) + transcription_model = Silero() audio_path = os.path.join("test_samples", "audio.wav") - transcription = silero.transcribe(audio_path) + transcription = transcription_model.transcribe(audio_path) assert similarity(TEXT, transcription) > MIN_SYNTHESIS_SCORE diff --git a/tests/test_synthesis.py b/tests/test_synthesis.py index 40c57e2..41f7de1 100644 --- a/tests/test_synthesis.py +++ b/tests/test_synthesis.py @@ -4,7 +4,7 @@ import librosa import torch from dataset.utils import similarity -from dataset.transcribe import create_transcription_model +from dataset.transcribe import Silero from synthesis.synthesize import load_model, synthesize from synthesis.vocoders import Hifigan from synthesis.vocoders.vocoder import Vocoder @@ -109,7 +109,7 @@ def test_hifigan_synthesis(): hifigan_model_path = os.path.join("test_samples", "hifigan.pt") hifigan_config_path = os.path.join("test_samples", "config.json") audio_path = "synthesized_audio.wav" - transcription_model = create_transcription_model() + transcription_model = Silero() hifigan = Hifigan(hifigan_model_path, hifigan_config_path) text = "the monkeys live"