german and spanish

app.py

@@ -1,18 +1,31 @@
 import os, sys, re
 import shutil
-import argparse
 import subprocess
 import soundfile
 from process_audio import segment_audio
 from write_srt import write_to_file
-from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC, Wav2Vec2Tokenizer
+from clean_text import clean_english, clean_german, clean_spanish
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
 import torch
 import gradio as gr
 
 
-model = "facebook/wav2vec2-large-960h-lv60-self"
-tokenizer = Wav2Vec2Tokenizer.from_pretrained(model)
-asr_model = Wav2Vec2ForCTC.from_pretrained(model)#.to('cuda')
+english_model = "facebook/wav2vec2-large-960h-lv60-self"
+english_tokenizer = Wav2Vec2Processor.from_pretrained(english_model)
+english_asr_model = Wav2Vec2ForCTC.from_pretrained(english_model)
+
+german_model = "jonatasgrosman/wav2vec2-large-xlsr-53-german"
+german_tokenizer = Wav2Vec2Processor.from_pretrained(german_model)
+german_asr_model = Wav2Vec2ForCTC.from_pretrained(german_model)
+
+spanish_model = "jonatasgrosman/wav2vec2-large-xlsr-53-spanish"
+spanish_tokenizer = Wav2Vec2Processor.from_pretrained(spanish_model)
+spanish_asr_model = Wav2Vec2ForCTC.from_pretrained(spanish_model)
+
+# Get German corpus and update nltk
+command = ["python", "-m", "textblob.download_corpora"]
+subprocess.run(command)
+
 
 # Line count for SRT file
 line_count = 0
@@ -34,18 +47,23 @@ def transcribe_audio(tokenizer, asr_model, audio_file, file_handle):
 
     
     infered_text = tokenizer.batch_decode(prediction)[0].lower()
-    infered_text = re.sub(r'  ', ' ', infered_text)
-    infered_text = re.sub(r'\bi\s', 'I ', infered_text)
-    infered_text = re.sub(r'\si$', ' I', infered_text)
-    infered_text = re.sub(r'i\'', 'I\'', infered_text)
-
-    limits = audio_file.split(os.sep)[-1][:-4].split("_")[-1].split("-")
-
     if len(infered_text) > 1:
+        if lang == 'english':
+            infered_text = clean_english(infered_text)
+        elif lang == 'german':
+            infered_text = clean_german(infered_text)
+        elif lang == 'spanish':
+            infered_text = clean_spanish(infered_text)
+
+        print(infered_text)
+        limits = audio_file.split(os.sep)[-1][:-4].split("_")[-1].split("-")
         line_count += 1
         write_to_file(file_handle, infered_text, line_count, limits)
+    else:
+        infered_text = ''
+
         
-def get_subs(input_file):
+def get_subs(input_file, language):
     # Get directory for audio
     base_directory = os.getcwd()
     audio_directory = os.path.join(base_directory, "audio")
@@ -71,6 +89,11 @@ def get_subs(input_file):
     file_handle.seek(0)
     for file in sort_alphanumeric(os.listdir(audio_directory)):
         audio_segment_path = os.path.join(audio_directory, file)
+        global lang
+        lang = language.lower()
+        tokenizer = globals()[lang+'_tokenizer']
+        asr_model = globals()[lang+'_asr_model']
+
         if audio_segment_path.split(os.sep)[-1] != audio_file.split(os.sep)[-1]:
             transcribe_audio(tokenizer, asr_model, audio_segment_path, file_handle)
 
@@ -84,7 +107,8 @@ gradio_ui = gr.Interface(
     fn=get_subs,
     title="Video to Subtitle",
     description="Get subtitles (SRT file) for your videos. Inference speed is about 10s/per 1min of video BUT the speed of uploading your video depends on your internet connection.",
-    inputs=gr.inputs.Video(label="Upload Video File"),
+    inputs=[gr.inputs.Video(label="Upload Video File"),
+        gr.inputs.Radio(label="Choose Language", choices=['English', 'German', 'Spanish'])],
     outputs=gr.outputs.File(label="Auto-Transcript")
     )
 

clean_text.py

@@ -0,0 +1,39 @@
+import os, re, string
+import subprocess
+from textblob_de import TextBlobDE as TextBlob
+
+
+def clean_english(text):
+    clean_text = re.sub(r'  ', ' ', text)
+    clean_text = re.sub(r'\bi\s', 'I ', clean_text)
+    clean_text = re.sub(r'\si$', ' I', clean_text)
+    clean_text = re.sub(r'i\'', 'I\'', clean_text)
+
+    return clean_text
+
+def clean_german(text):
+	text = text.translate(str.maketrans('', '', string.punctuation))
+	
+	# Tokenize German text
+	blob = TextBlob(text)
+	pos = blob.tags
+
+	# Get nouns and capitalize 
+	nouns = {}
+	for idx in pos:
+		if idx[1] == 'NN' and len(idx[0]) > 1:
+			nouns[idx[0]] = idx[0].capitalize()
+		
+	if len(nouns) != 0:
+		pattern = re.compile("|".join(nouns.keys()))
+		text = pattern.sub(lambda m: nouns[re.escape(m.group(0))], text)
+	
+	return text
+
+
+def clean_spanish(text):
+	clean_text = text.translate(str.maketrans('', '', string.punctuation))
+	clean_text = re.sub(r'  ', ' ', clean_text)
+
+	return clean_text
+

requirements.txt

@@ -3,3 +3,4 @@ transformers
 torch
 gradio
 auditok
+textblob_de