salmanmapkar's picture
Update app.py
eea7953
from __future__ import unicode_literals
import youtube_dl
import yt_dlp
from pydub import AudioSegment
from pyannote.audio import Pipeline
import re
import whisper
import os
import ffmpeg
import subprocess
import gradio as gr
import traceback
import json
pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token="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")
from pydub.effects import speedup
import moviepy.editor as mp
import datetime
import torch
import pyannote.audio
from pyannote.audio.pipelines.speaker_verification import SpeechBrainPretrainedSpeakerEmbedding #PyannoteAudioPretrainedSpeakerEmbedding
from pyannote.audio import Audio
from pyannote.core import Segment
import wave
import contextlib
from sklearn.cluster import AgglomerativeClustering
import numpy as np
import json
from datetime import timedelta
from transformers import T5ForConditionalGeneration, T5Tokenizer
__FILES = set()
wispher_models = list(whisper._MODELS.keys())
def correct_grammar(input_text,num_return_sequences=1):
torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'
tokenizer = T5Tokenizer.from_pretrained('deep-learning-analytics/GrammarCorrector')
model = T5ForConditionalGeneration.from_pretrained('deep-learning-analytics/GrammarCorrector').to(torch_device)
batch = tokenizer([input_text],truncation=True,padding='max_length',max_length=len(input_text), return_tensors="pt").to(torch_device)
results = model.generate(**batch,max_length=len(input_text),num_beams=2, num_return_sequences=num_return_sequences, temperature=1.5)
generated_sequences = []
for generated_sequence_idx, generated_sequence in enumerate(results):
text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True, skip_special_tokens=True)
generated_sequences.append(text)
generated_text = "".join(generated_sequences)
_generated_text = ""
for idx, _sentence in enumerate(generated_text.split('.'), 0):
if not idx:
_generated_text+=_sentence+'.'
elif _sentence[:1]!=' ':
_generated_text+=' '+_sentence+'.'
elif _sentence[:1]=='':
pass
else:
_generated_text+=_sentence+'.'
return _generated_text
def CreateFile(filename):
__FILES.add(filename)
return filename
def RemoveFile(filename):
if (os.path.isfile(filename)):
os.remove(filename)
def RemoveAllFiles():
for file in __FILES:
if (os.path.isfile(file)):
os.remove(file)
def Transcribe_V1(NumberOfSpeakers, SpeakerNames="", audio="temp_audio.wav"):
SPEAKER_DICT = {}
SPEAKERS = [speaker.strip() for speaker in SpeakerNames.split(',') if len(speaker)]
def GetSpeaker(sp):
speaker = sp
if sp not in list(SPEAKER_DICT.keys()):
if len(SPEAKERS):
t = SPEAKERS.pop(0)
SPEAKER_DICT[sp] = t
speaker = SPEAKER_DICT[sp]
else:
speaker = SPEAKER_DICT[sp]
return speaker
def millisec(timeStr):
spl = timeStr.split(":")
s = (int)((int(spl[0]) * 60 * 60 + int(spl[1]) * 60 + float(spl[2]) )* 1000)
return s
def preprocess(audio):
t1 = 0 * 1000
t2 = 20 * 60 * 1000
newAudio = AudioSegment.from_wav(audio)
a = newAudio[t1:t2]
spacermilli = 2000
spacer = AudioSegment.silent(duration=spacermilli)
newAudio = spacer.append(a, crossfade=0)
newAudio.export(audio, format="wav")
return spacermilli, spacer
def diarization(audio):
as_audio = AudioSegment.from_wav(audio)
DEMO_FILE = {'uri': 'blabal', 'audio': audio}
if NumberOfSpeakers:
dz = pipeline(DEMO_FILE, num_speakers=NumberOfSpeakers)
else:
dz = pipeline(DEMO_FILE)
with open(CreateFile(f"diarization_{audio}.txt"), "w") as text_file:
text_file.write(str(dz))
dz = open(CreateFile(f"diarization_{audio}.txt")).read().splitlines()
dzList = []
for l in dz:
start, end = tuple(re.findall('[0-9]+:[0-9]+:[0-9]+\.[0-9]+', string=l))
start = millisec(start)
end = millisec(end)
lex = GetSpeaker(re.findall('(SPEAKER_[0-9][0-9])', string=l)[0])
dzList.append([start, end, lex])
sounds = spacer
segments = []
dz = open(CreateFile(f"diarization_{audio}.txt")).read().splitlines()
for l in dz:
start, end = tuple(re.findall('[0-9]+:[0-9]+:[0-9]+\.[0-9]+', string=l))
start = millisec(start)
end = millisec(end)
segments.append(len(sounds))
sounds = sounds.append(as_audio[start:end], crossfade=0)
sounds = sounds.append(spacer, crossfade=0)
sounds.export(CreateFile(f"dz_{audio}.wav"), format="wav")
return f"dz_{audio}.wav", dzList, segments
def transcribe(dz_audio):
model = whisper.load_model("medium")
result = model.transcribe(dz_audio)
# for _ in result['segments']:
# print(_['start'], _['end'], _['text'])
captions = [[((caption["start"]*1000)), ((caption["end"]*1000)), caption["text"]] for caption in result['segments']]
conversation = []
for i in range(len(segments)):
idx = 0
for idx in range(len(captions)):
if captions[idx][0] >= (segments[i] - spacermilli):
break;
while (idx < (len(captions))) and ((i == len(segments) - 1) or (captions[idx][1] < segments[i+1])):
c = captions[idx]
start = dzList[i][0] + (c[0] -segments[i])
if start < 0:
start = 0
idx += 1
if not len(conversation):
conversation.append([dzList[i][2], c[2]])
elif conversation[-1][0] == dzList[i][2]:
conversation[-1][1] += c[2]
else:
conversation.append([dzList[i][2], c[2]])
#print(f"[{dzList[i][2]}] {c[2]}")
return conversation, ("".join([f"{speaker} --> {text}\n" for speaker, text in conversation]))
spacermilli, spacer = preprocess(audio)
dz_audio, dzList, segments = diarization(audio)
conversation, t_text = transcribe(dz_audio)
RemoveAllFiles()
return (t_text, ({ "data": [{"speaker": speaker, "text": text} for speaker, text in conversation]}))
def Transcribe_V2(model, num_speakers, speaker_names, audio="temp_audio.wav"):
model = whisper.load_model(model)
# embedding_model = SpeechBrainPretrainedSpeakerEmbedding("speechbrain/spkrec-ecapa-voxceleb")
embedding_model = SpeechBrainPretrainedSpeakerEmbedding(
"speechbrain/spkrec-ecapa-voxceleb",
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
)
SPEAKER_DICT = {}
default_speaker_names = ['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']
SPEAKERS = [speaker.strip() for speaker in speaker_names.split(',') if len(speaker)]
def GetSpeaker(sp):
speaker = sp
if sp not in list(SPEAKER_DICT.keys()):
if len(SPEAKERS):
t = SPEAKERS.pop(0)
SPEAKER_DICT[sp] = t
speaker = SPEAKER_DICT[sp]
elif len(default_speaker_names):
t = default_speaker_names.pop(0)
SPEAKER_DICT[sp] = t
speaker = SPEAKER_DICT[sp]
else:
speaker = SPEAKER_DICT[sp]
return speaker
# audio = Audio()
def diarization(audio):
def millisec(timeStr):
spl = timeStr.split(":")
s = (int)((int(spl[0]) * 60 * 60 + int(spl[1]) * 60 + float(spl[2]) )* 1000)
return s
as_audio = AudioSegment.from_wav(audio)
DEMO_FILE = {'uri': 'blabal', 'audio': audio}
hparams = pipeline.parameters(instantiated=True)
hparams["segmentation"]["min_duration_off"] -= 0.25
pipeline.instantiate(hparams)
if num_speakers:
dz = pipeline(DEMO_FILE, num_speakers=num_speakers)
else:
dz = pipeline(DEMO_FILE)
with open(CreateFile(f"diarization_{audio}.txt"), "w") as text_file:
text_file.write(str(dz))
dz = open(CreateFile(f"diarization_{audio}.txt")).read().splitlines()
print(dz)
dzList = []
for l in dz:
start, end = tuple(re.findall('[0-9]+:[0-9]+:[0-9]+\.[0-9]+', string=l))
start = millisec(start)
end = millisec(end)
lex = GetSpeaker(re.findall('(SPEAKER_[0-9][0-9])', string=l)[0])
dzList.append([start, end, lex])
return dzList
def get_output(segments):
# print(segments)
conversation=[]
for (i, segment) in enumerate(segments):
# print(f"{i}, {segment["speaker"]}, {segments[i - 1]["speaker"]}, {}")
if not len(conversation):
conversation.append([str(timedelta(seconds=float(segment['start']))),str(timedelta(seconds=float(segment['end']))),GetSpeaker(segment["speaker"]), segment["text"].lstrip()])
elif conversation[-1][2] == GetSpeaker(segment["speaker"]):
conversation[-1][3] += segment["text"].lstrip()
else:
conversation.append([str(timedelta(seconds=float(segment['start']))),str(timedelta(seconds=float(segment['end']))),GetSpeaker(segment["speaker"]), segment["text"].lstrip()])
# if i == 0 or segments[i - 1]["speaker"] != segment["speaker"]:
# if i != 0:
# conversation.append([GetSpeaker(segment["speaker"]), segment["text"][1:]]) # segment["speaker"] + ' ' + str(time(segment["start"])) + '\n\n'
# conversation[-1][1] += segment["text"][1:]
# return output
for idx in range(len(conversation)):
conversation[idx][3] = correct_grammar(conversation[idx][3])
return ("".join([f"[{start}] - {speaker} \n{text}\n" for start, end, speaker, text in conversation])), ({ "data": [{"start": start, "end":end, "speaker": speaker, "text": text} for start, end, speaker, text in conversation]})
def get_duration(path):
with contextlib.closing(wave.open(path,'r')) as f:
frames = f.getnframes()
rate = f.getframerate()
return frames / float(rate)
def make_embeddings(path, segments, duration):
embeddings = np.zeros(shape=(len(segments), 192))
for i, segment in enumerate(segments):
embeddings[i] = segment_embedding(path, segment, duration)
return np.nan_to_num(embeddings)
def segment_embedding(path, segment, duration):
start = segment["start"]
# Whisper overshoots the end timestamp in the last segment
end = min(duration, segment["end"])
clip = Segment(start, end)
waveform, sample_rate = Audio().crop(path, clip)
return embedding_model(waveform[None])
def add_speaker_labels(segments, embeddings, num_speakers):
clustering = AgglomerativeClustering(num_speakers).fit(embeddings)
labels = clustering.labels_
for i in range(len(segments)):
segments[i]["speaker"] = 'SPEAKER ' + str(labels[i] + 1)
def time(secs):
return datetime.timedelta(seconds=round(secs))
duration = get_duration(audio)
if duration > 4 * 60 * 60:
return "Audio duration too long"
# print(json.dumps(diarization(audio)))
result = model.transcribe(audio)
# print(json.dumps(result))
segments = result["segments"]
num_speakers = min(max(round(num_speakers), 1), len(segments))
if len(segments) == 1:
segments[0]['speaker'] = 'SPEAKER 1'
else:
embeddings = make_embeddings(audio, segments, duration)
add_speaker_labels(segments, embeddings, num_speakers)
return get_output(segments)
# return output
def AudioTranscribe(NumberOfSpeakers=None, SpeakerNames="", audio="", retries=5, model='base'):
print(f"{NumberOfSpeakers}, {SpeakerNames}, {retries}")
if retries:
# subprocess.call(['ffmpeg', '-i', audio,'temp_audio.wav'])
try:
subprocess.call(['ffmpeg', '-i', audio,'temp_audio.wav'])
except Exception as ex:
traceback.print_exc()
return AudioTranscribe(NumberOfSpeakers, SpeakerNames, audio, retries-1)
if not (os.path.isfile("temp_audio.wav")):
return AudioTranscribe(NumberOfSpeakers, SpeakerNames, audio, retries-1)
return Transcribe_V2(model, NumberOfSpeakers, SpeakerNames)
else:
raise gr.Error("There is some issue ith Audio Transcriber. Please try again later!")
def VideoTranscribe(NumberOfSpeakers=None, SpeakerNames="", video="", retries=5, model='base'):
if retries:
try:
clip = mp.VideoFileClip(video)
clip.audio.write_audiofile("temp_audio.wav")
# command = f"ffmpeg -i {video} -ab 160k -ac 2 -ar 44100 -vn temp_audio.wav"
# subprocess.call(command, shell=True)
except Exception as ex:
traceback.print_exc()
return VideoTranscribe(NumberOfSpeakers, SpeakerNames, video, retries-1)
if not (os.path.isfile("temp_audio.wav")):
return VideoTranscribe(NumberOfSpeakers, SpeakerNames, video, retries-1)
return Transcribe_V2(model, NumberOfSpeakers, SpeakerNames)
else:
raise gr.Error("There is some issue ith Video Transcriber. Please try again later!")
def YoutubeTranscribe(NumberOfSpeakers=None, SpeakerNames="", URL="", retries = 5, model='base'):
if retries:
if "youtu" not in URL.lower():
raise gr.Error(f"{URL} is not a valid youtube URL.")
else:
RemoveFile("temp_audio.wav")
ydl_opts = {
'format': 'bestaudio/best',
'outtmpl': 'temp_audio.%(ext)s',
'postprocessors': [{
'key': 'FFmpegExtractAudio',
'preferredcodec': 'wav',
}],
}
try:
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
ydl.download([URL])
except:
return YoutubeTranscribe(NumberOfSpeakers, SpeakerNames, URL, retries-1)
stream = ffmpeg.input('temp_audio.m4a')
stream = ffmpeg.output(stream, 'temp_audio.wav')
RemoveFile("temp_audio.m4a")
return Transcribe_V2(model, NumberOfSpeakers, SpeakerNames)
else:
raise gr.Error(f"Unable to get video from {URL}")
with gr.Blocks() as yav_ui:
with gr.Row():
with gr.Column():
with gr.Tab("Youtube", id=1):
ysz = gr.Dropdown(label="Model Size", choices=wispher_models , value='base')
yinput_nos = gr.Number(label="Number of Speakers", placeholder="2")
yinput_sn = gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2")
yinput = gr.Textbox(label="Youtube Link", placeholder="https://www.youtube.com/watch?v=GECcjrYHH8w")
ybutton_transcribe = gr.Button("Transcribe", show_progress=True, scroll_to_output=True)
with gr.Tab("Video", id=2):
vsz = gr.Dropdown(label="Model Size", choices=wispher_models, value='base')
vinput_nos = gr.Number(label="Number of Speakers", placeholder="2")
vinput_sn = gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2")
vinput = gr.Video(label="Video")
vbutton_transcribe = gr.Button("Transcribe", show_progress=True, scroll_to_output=True)
with gr.Tab("Audio", id=3):
asz = gr.Dropdown(label="Model Size", choices=wispher_models , value='base')
ainput_nos = gr.Number(label="Number of Speakers", placeholder="2")
ainput_sn = gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2")
ainput = gr.Audio(label="Audio", type="filepath")
abutton_transcribe = gr.Button("Transcribe", show_progress=True, scroll_to_output=True)
with gr.Column():
with gr.Tab("Text"):
output_textbox = gr.Textbox(label="Transcribed Text", lines=15)
with gr.Tab("JSON"):
output_json = gr.JSON(label="Transcribed JSON")
ybutton_transcribe.click(
fn=YoutubeTranscribe,
inputs=[yinput_nos,yinput_sn,yinput, ysz],
outputs=[output_textbox,output_json]
)
abutton_transcribe.click(
fn=AudioTranscribe,
inputs=[ainput_nos,ainput_sn,ainput, asz],
outputs=[output_textbox,output_json]
)
vbutton_transcribe.click(
fn=VideoTranscribe,
inputs=[vinput_nos,vinput_sn,vinput, vsz],
outputs=[output_textbox,output_json]
)
yav_ui.launch(debug=True)