|
import gradio as gr |
|
from transformers import pipeline, Wav2Vec2ProcessorWithLM, Wav2Vec2ForCTC |
|
import os |
|
import soundfile as sf |
|
from pyannote.audio import Pipeline |
|
import torch |
|
from pydub import AudioSegment |
|
from pydub.playback import play |
|
from datetime import datetime, timedelta |
|
import time |
|
from silero_vad import load_silero_vad, read_audio, get_speech_timestamps |
|
|
|
HF_TOKEN = os.environ.get("HF_TOKEN") |
|
sr = 16000 |
|
channels = 1 |
|
|
|
model_name = "Mihaj/wav2vec2-large-xls-r-300m-ruOH-alphav" |
|
bond005_model = "bond005/wav2vec2-large-ru-golos-with-lm" |
|
processor = Wav2Vec2ProcessorWithLM.from_pretrained(bond005_model) |
|
model = Wav2Vec2ForCTC.from_pretrained(bond005_model) |
|
pipe = pipeline("automatic-speech-recognition", model=model, tokenizer=processor, feature_extractor=processor.feature_extractor, decoder=processor.decoder) |
|
model = load_silero_vad() |
|
|
|
pipeline_dia = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", |
|
use_auth_token=HF_TOKEN) |
|
|
|
|
|
temp_path = "temp.wav" |
|
|
|
def preprocess(audio_path): |
|
print("PREPROCESSING STARTED") |
|
sound = AudioSegment.from_file(audio_path, format="mp3") |
|
sound = sound.set_frame_rate(sr) |
|
sound = sound.set_channels(channels) |
|
sound.export(temp_path, format="wav") |
|
print("PREPROCESSING ENDED") |
|
return temp_path |
|
|
|
def fast_transcribe(diarise, how_diarise, translate, audio): |
|
audio = preprocess(audio) |
|
y, sr = sf.read(audio) |
|
if diarise: |
|
if how_diarise=="Accurate": |
|
print("DIARISING") |
|
dia = pipeline_dia(audio) |
|
print("DIARISING ENDED") |
|
lines = [] |
|
for i, line in enumerate(dia.to_lab().split('\n')): |
|
if line.strip() != "": |
|
res = line.split(" ") |
|
start = int(float(res[0]) * sr) |
|
start_time = str(datetime.fromtimestamp(start / sr) - timedelta(hours=1, minutes=0)).split()[1] |
|
start_time_prts = start_time.split(":") |
|
start_time_srt = f"{start_time_prts[0]}:{start_time_prts[1]}:{float(start_time_prts[2]):.3f}".replace('.', ',') |
|
end = int(float(res[1]) * sr) |
|
end_time = str(datetime.fromtimestamp(end / sr) - timedelta(hours=1, minutes=0)).split()[1] |
|
end_time_prts = end_time.split(":") |
|
end_time_srt = f"{end_time_prts[0]}:{end_time_prts[1]}:{float(end_time_prts[2]):.3f}".replace('.', ',') |
|
label = res[2] |
|
print(f"RECOGNISING LINE_{i} T_START {start_time_srt} T_END {end_time_srt} SPEAKER_{label}") |
|
trans = pipe(y[start:end], chunk_length_s=10, stride_length_s=(4, 2))["text"] |
|
if not translate: |
|
lines.append(f"{i+1}\n{start_time_srt} --> {end_time_srt}\n[{label}] {trans}\n") |
|
else: |
|
print("TRANSLATION STARTED") |
|
trans_eng = translator.translate('trans', src='ru', dest="en").text |
|
print(f"TRANSLATION ENDED RESULT {trans_eng}") |
|
lines.append(f"{i+1}\n{start_time_srt} --> {end_time_srt}\n[{label}] {trans}\n[{label}] {trans_eng}\n") |
|
print("RECOGNISING ENDED") |
|
print(f"LINE RESULT {trans}") |
|
else: |
|
print("DIARISING") |
|
wav = read_audio(audio) |
|
speech_timestamps = get_speech_timestamps(wav, model, speech_pad_ms=80, min_silence_duration_ms=150, window_size_samples=256) |
|
print("DIARISING ENDED") |
|
lines = [] |
|
for i, line in enumerate(speech_timestamps): |
|
start = line['start'] |
|
print(start) |
|
start_time = str(datetime.fromtimestamp(start / sr) - timedelta(hours=1, minutes=0)).split()[1] |
|
start_time_prts = start_time.split(":") |
|
start_time_srt = f"{start_time_prts[0]}:{start_time_prts[1]}:{float(start_time_prts[2]):.3f}".replace('.', ',') |
|
print(start_time_srt) |
|
end = line['end'] |
|
end_time = str(datetime.fromtimestamp(end / sr) - timedelta(hours=1, minutes=0)).split()[1] |
|
end_time_prts = end_time.split(":") |
|
end_time_srt = f"{end_time_prts[0]}:{end_time_prts[1]}:{float(end_time_prts[2]):.3f}".replace('.', ',') |
|
print(f"RECOGNISING LINE_{i} T_START {start_time_srt} T_END {end_time_srt}") |
|
trans = pipe(y[start:end], chunk_length_s=10, stride_length_s=(4, 2))["text"] |
|
print("RECOGNISING ENDED") |
|
if not translate: |
|
lines.append(f"{i+1}\n{start_time_srt} --> {end_time_srt}\n[{trans}\n") |
|
else: |
|
print("TRANSLATION STARTED") |
|
trans_eng = translator.translate(trans, src='ru', dest="en").text |
|
print(f"TRANSLATION ENDED RESULT {trans_eng}") |
|
lines.append(f"{i+1}\n{start_time_srt} --> {end_time_srt}\n{trans}\n{trans_eng}\n") |
|
|
|
print(f"LINE RESULT {trans}") |
|
text = "\n".join(lines) |
|
else: |
|
print("RECOGNISING FULL AUDIO") |
|
res = pipe(y, chunk_length_s=10, stride_length_s=(4, 2)) |
|
print("RECOGNISING FULL AUDIO ENDED") |
|
text = res["text"] |
|
return text |
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown(""" |
|
#Wav2Vec2 RuOH |
|
Realtime demo for Russian Oral History recognition using several diarizations method (Silero VAD, Pyannote) and a Wav2Vec large model from bond005. https://huggingface.co/bond005/wav2vec2-large-ru-golos-with-lm" |
|
""") |
|
with gr.Tab("Fast Translation"): |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
fast_diarize_input = gr.Checkbox(label="Subtitles", info="Do you want subtitles?") |
|
fast_diarize_radio_input = gr.Radio(["Fast", "Accurate", "-"], label="separating_on_subtitles_pption", info="You can choose separating audio on smaller pieces by faster yet low quality variant (Silero VAD), or slower yet high quality variant (Pyannote.Diarization, this option will detect different speakers)") |
|
fast_translate_input = gr.Checkbox(label="Translate", info="Do you want translation to English?") |
|
fast_audio_input = gr.Audio(type="filepath") |
|
|
|
fast_output = gr.Textbox() |
|
|
|
fast_inputs = [fast_diarize_input, fast_diarize_radio_input, fast_translate_input, fast_audio_input] |
|
fast_recognize_button = gr.Button("Run") |
|
|
|
|
|
fast_recognize_button.click(fast_transcribe, inputs=fast_inputs, outputs=fast_output) |
|
|
|
if __name__ == "__main__": |
|
demo.launch() |
