app.py

# Inspiration from https://huggingface.co/spaces/vumichien/whisper-speaker-diarization

import whisper
import datetime
import subprocess
import gradio as gr
from pathlib import Path
import pandas as pd
import re
import time
import os 
import numpy as np
from sklearn.cluster import AgglomerativeClustering

from pytube import YouTube
import torch
import pyannote.audio
from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding
from pyannote.audio import Audio
from pyannote.core import Segment

from gpuinfo import GPUInfo

import wave
import contextlib
from transformers import pipeline
import psutil

from zipfile import ZipFile
from io import StringIO
import csv

# ---- Model Loading ----

whisper_models = ["base", "small", "medium", "large"]
source_languages = {
    "en": "English",
    "de": "German",
    "es": "Spanish",
    "fr": "French",
}

source_language_list = [key[0] for key in source_languages.items()]

MODEL_NAME = "openai/whisper-small"
lang = "en"

device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = pipeline(
    task="automatic-speech-recognition",
    model=MODEL_NAME,
    chunk_length_s=30,
    device=device,
)

pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language=lang, task="transcribe")

embedding_model = PretrainedSpeakerEmbedding( 
    "speechbrain/spkrec-ecapa-voxceleb",
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"))

# ---- S2T & Speaker diarization ----

def transcribe(microphone, file_upload):
    warn_output = ""
    if (microphone is not None) and (file_upload is not None):
        warn_output = (
            "WARNING: You've uploaded an audio file and used the microphone. "
            "The recorded file from the microphone will be used and the uploaded audio will be discarded.\n"
        )

    elif (microphone is None) and (file_upload is None):
        return "ERROR: You have to either use the microphone or upload an audio file"

    file = microphone if microphone is not None else file_upload

    text = pipe(file)["text"]

    return warn_output + text


def convert_time(secs):
    return datetime.timedelta(seconds=round(secs))

def convert_to_wav(filepath):
    _,file_ending = os.path.splitext(f'{filepath}')
    audio_file = filepath.replace(file_ending, ".wav")
    print("starting conversion to wav")
    os.system(f'ffmpeg -i "{filepath}" -ar 16000 -ac 1 -c:a pcm_s16le "{audio_file}"')
    return audio_file


def speech_to_text(microphone, file_upload, selected_source_lang, whisper_model, num_speakers):
    """
    # Transcribe audio file and separate into segment, assign speakers to segments
    1. Using Open AI's Whisper model to seperate audio into segments and generate transcripts.
    2. Generating speaker embeddings for each segments.
    3. Applying agglomerative clustering on the embeddings to identify the speaker for each segment.
    
    Speech Recognition is based on models from OpenAI Whisper https://github.com/openai/whisper
    Speaker diarization model and pipeline from by https://github.com/pyannote/pyannote-audio
    """
    
    model = whisper.load_model(whisper_model)
    time_start = time.time()

    try:
        # Read and convert audio file
        warn_output = ""
        if (microphone is not None) and (file_upload is not None):
            warn_output = (
                "WARNING: You've uploaded an audio file and used the microphone. "
                "The recorded file from the microphone will be used and the uploaded audio will be discarded.\n"
            )

        elif (microphone is None) and (file_upload is None):
            return "ERROR: You have to either use the microphone or upload an audio file"

        file = microphone if microphone is not None else file_upload
        
        if microphone is None and file_upload is not None:
            file = convert_to_wav(file)
        
        # Get duration
        with contextlib.closing(wave.open(file,'r')) as f:
            frames = f.getnframes()
            rate = f.getframerate()
            duration = frames / float(rate)
        print(f"conversion to wav ready, duration of audio file: {duration}")

        # Transcribe audio
        options = dict(language=selected_source_lang, beam_size=3, best_of=3)
        transcribe_options = dict(task="transcribe", **options)
        result = model.transcribe(file, **transcribe_options)
        segments = result["segments"]
        print("whisper done with transcription")
    except Exception as e:
        raise RuntimeError("Error converting audio file")

    try:
        # Create embedding
        def segment_embedding(segment):
            audio = Audio()
            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(file, clip)
            return embedding_model(waveform[None])

        embeddings = np.zeros(shape=(len(segments), 192))
        for i, segment in enumerate(segments):
            embeddings[i] = segment_embedding(segment)
        embeddings = np.nan_to_num(embeddings)
        print(f'Embedding shape: {embeddings.shape}')

        # Assign speaker label
        if num_speakers == 1:
            for i in range(len(segments)):
                segments[i]["speaker"] = 'SPEAKER 1'
        else:
            clustering = AgglomerativeClustering(num_speakers).fit(embeddings)
            labels = clustering.labels_
            for i in range(len(segments)):
                segments[i]["speaker"] = 'SPEAKER ' + str(labels[i] + 1)

        # Make output
        objects = {
            'Start' : [],
            'End': [],
            'Speaker': [],
            'Text': []
        }
        text = ''
        if num_speakers == 1:
            objects['Start'].append(str(convert_time(segment["start"])))
            objects['Speaker'].append(segment["speaker"])
            for (i, segment) in enumerate(segments):
                text += segment["text"] + ' '
            objects['Text'].append(text)
            objects['End'].append(str(convert_time(segment["end"])))
        else:
            for (i, segment) in enumerate(segments):
                if i == 0 or segments[i - 1]["speaker"] != segment["speaker"]:
                    objects['Start'].append(str(convert_time(segment["start"])))
                    objects['Speaker'].append(segment["speaker"])
                    if i != 0:
                        objects['End'].append(str(convert_time(segments[i - 1]["end"])))
                        objects['Text'].append(text)
                        text = ''
                text += segment["text"] + ' '
            objects['End'].append(str(convert_time(segments[i - 1]["end"])))
            objects['Text'].append(text)
        
        time_end = time.time()
        time_diff = time_end - time_start
        memory = psutil.virtual_memory()
        gpu_utilization, gpu_memory = GPUInfo.gpu_usage()
        gpu_utilization = gpu_utilization[0] if len(gpu_utilization) > 0 else 0
        gpu_memory = gpu_memory[0] if len(gpu_memory) > 0 else 0
        system_info = f"""
        *Memory: {memory.total / (1024 * 1024 * 1024):.2f}GB, used: {memory.percent}%, available: {memory.available / (1024 * 1024 * 1024):.2f}GB.* 
        *Processing time: {time_diff:.5} seconds.*
        *GPU Utilization: {gpu_utilization}%, GPU Memory: {gpu_memory}MiB.*
        """

        return pd.DataFrame(objects), system_info
    
    except Exception as e:
        raise RuntimeError("Error Running inference with local model", e)

# ---- Youtube Conversion ----

def get_youtube(video_url):
    yt = YouTube(video_url)
    abs_video_path = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().first().download()
    print("Success download video")
    print(abs_video_path)
    return abs_video_path



def yt_to_text(video_file_path, selected_source_lang, whisper_model, num_speakers):
    """
    # Transcribe youtube link using OpenAI Whisper
    1. Using Open AI's Whisper model to seperate audio into segments and generate transcripts.
    2. Generating speaker embeddings for each segments.
    3. Applying agglomerative clustering on the embeddings to identify the speaker for each segment.
    
    Speech Recognition is based on models from OpenAI Whisper https://github.com/openai/whisper
    Speaker diarization model and pipeline from by https://github.com/pyannote/pyannote-audio
    """
    
    model = whisper.load_model(whisper_model)
    time_start = time.time()
    if(video_file_path == None):
        raise ValueError("Error no video input")
    print(video_file_path)

    try:
        # Read and convert youtube video
        _,file_ending = os.path.splitext(f'{video_file_path}')
        print(f'file ending is {file_ending}')
        audio_file = video_file_path.replace(file_ending, ".wav")
        print("starting conversion to wav")
        os.system(f'ffmpeg -i "{video_file_path}" -ar 16000 -ac 1 -c:a pcm_s16le "{audio_file}"')
        
        # Get duration
        with contextlib.closing(wave.open(audio_file,'r')) as f:
            frames = f.getnframes()
            rate = f.getframerate()
            duration = frames / float(rate)
        print(f"conversion to wav ready, duration of audio file: {duration}")

        # Transcribe audio
        options = dict(language=selected_source_lang, beam_size=5, best_of=5)
        transcribe_options = dict(task="transcribe", **options)
        result = model.transcribe(audio_file, **transcribe_options)
        segments = result["segments"]
        print("starting whisper done with whisper")
    except Exception as e:
        raise RuntimeError("Error converting video to audio")

    try:
        # Create embedding
        def segment_embedding(segment):
            audio = Audio()
            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(audio_file, clip)
            return embedding_model(waveform[None])

        embeddings = np.zeros(shape=(len(segments), 192))
        for i, segment in enumerate(segments):
            embeddings[i] = segment_embedding(segment)
        embeddings = np.nan_to_num(embeddings)
        print(f'Embedding shape: {embeddings.shape}')

        # Assign speaker label
        if num_speakers == 1:
            for i in range(len(segments)):
                segments[i]["speaker"] = 'SPEAKER 1'
        else:
            clustering = AgglomerativeClustering(num_speakers).fit(embeddings)
            labels = clustering.labels_
            for i in range(len(segments)):
                segments[i]["speaker"] = 'SPEAKER ' + str(labels[i] + 1)

        # Make output
        objects = {
            'Start' : [],
            'End': [],
            'Speaker': [],
            'Text': []
        }
        text = ''
        if num_speakers == 1:
            objects['Start'].append(str(convert_time(segment["start"])))
            objects['Speaker'].append(segment["speaker"])
            for (i, segment) in enumerate(segments):
                text += segment["text"] + ' '
            objects['Text'].append(text)
            objects['End'].append(str(convert_time(segment["end"])))
        else:
            for (i, segment) in enumerate(segments):
                if i == 0 or segments[i - 1]["speaker"] != segment["speaker"]:
                    objects['Start'].append(str(convert_time(segment["start"])))
                    objects['Speaker'].append(segment["speaker"])
                    if i != 0:
                        objects['End'].append(str(convert_time(segments[i - 1]["end"])))
                        objects['Text'].append(text)
                        text = ''
                text += segment["text"] + ' '
            objects['End'].append(str(convert_time(segments[i - 1]["end"])))
            objects['Text'].append(text)
        
        time_end = time.time()
        time_diff = time_end - time_start
        memory = psutil.virtual_memory()
        gpu_utilization, gpu_memory = GPUInfo.gpu_usage()
        gpu_utilization = gpu_utilization[0] if len(gpu_utilization) > 0 else 0
        gpu_memory = gpu_memory[0] if len(gpu_memory) > 0 else 0
        system_info = f"""
        *Memory: {memory.total / (1024 * 1024 * 1024):.2f}GB, used: {memory.percent}%, available: {memory.available / (1024 * 1024 * 1024):.2f}GB.* 
        *Processing time: {time_diff:.5} seconds.*
        *GPU Utilization: {gpu_utilization}%, GPU Memory: {gpu_memory}MiB.*
        """

        return pd.DataFrame(objects), system_info
    
    except Exception as e:
        raise RuntimeError("Error Running inference with local model", e)

def download_csv(dataframe: pd.DataFrame):
    compression_options = dict(method='zip', archive_name='output.csv')
    dataframe.to_csv('output.zip', index=False, compression=compression_options)  
    return 'output.zip'

# ---- Gradio Layout ----
# Inspiration from https://huggingface.co/spaces/vumichien/whisper-speaker-diarization

# -- General Functions --
df_init = pd.DataFrame(columns=['Start', 'End', 'Speaker', 'Text'])
memory = psutil.virtual_memory()
title = "Whisper speaker diarization & speech recognition"
interface = gr.Blocks(title=title)
interface.encrypt = False

# -- Functions Audio Input --
microphone_in = gr.inputs.Audio(source="microphone", 
                                type="filepath", 
                                optional=True)

upload_in = gr.inputs.Audio(source="upload", 
                            type="filepath", 
                            optional=True)

selected_source_lang_audio = gr.Dropdown(choices=source_language_list, 
                                         type="value", 
                                         value="en", 
                                         label="Spoken language in audio", 
                                         interactive=True)

selected_whisper_model_audio = gr.Dropdown(choices=whisper_models, 
                                           type="value", 
                                           value="base", 
                                           label="Selected Whisper model", 
                                           interactive=True)

number_speakers_audio = gr.Number(precision=0, 
                                  value=2, 
                                  label="Selected number of speakers", 
                                  interactive=True)

system_info_audio = gr.Markdown(f"*Memory: {memory.total / (1024 * 1024 * 1024):.2f}GB, used: {memory.percent}%, available: {memory.available / (1024 * 1024 * 1024):.2f}GB*")

transcription_df_audio = gr.DataFrame(value=df_init,
                                      label="Transcription dataframe", 
                                      row_count=(0, "dynamic"), 
                                      max_rows = 10, 
                                      wrap=True, 
                                      overflow_row_behaviour='paginate')

csv_download_audio = gr.outputs.File(label="Download CSV")

# -- Functions Video Input --
video_in = gr.Video(label="Video file", 
                    mirror_webcam=False)

youtube_url_in = gr.Textbox(label="Youtube url", 
                            lines=1, 
                            interactive=True)

selected_source_lang_yt = gr.Dropdown(choices=source_language_list, 
                                      type="value", 
                                      value="en", 
                                      label="Spoken language in audio", 
                                      interactive=True)

selected_whisper_model_yt = gr.Dropdown(choices=whisper_models, 
                                        type="value", 
                                        value="base", 
                                        label="Selected Whisper model", 
                                        interactive=True)

number_speakers_yt = gr.Number(precision=0, 
                               value=2, 
                               label="Selected number of speakers", 
                               interactive=True)

system_info_yt = gr.Markdown(f"*Memory: {memory.total / (1024 * 1024 * 1024):.2f}GB, used: {memory.percent}%, available: {memory.available / (1024 * 1024 * 1024):.2f}GB*")

transcription_df_yt = gr.DataFrame(value=df_init,
                                   label="Transcription dataframe", 
                                   row_count=(0, "dynamic"), 
                                   max_rows = 10, 
                                   wrap=True, 
                                   overflow_row_behaviour='paginate')

csv_download_yt = gr.outputs.File(label="Download CSV")

with interface:
    with gr.Tab("Whisper speaker diarization & speech recognition"):
        gr.Markdown('''
            <div>
            <h1 style='text-align: center'>Whisper speaker diarization & speech recognition</h1>
            This space uses Whisper models from <a href='https://github.com/openai/whisper' target='_blank'><b>OpenAI</b></a> to recoginze the speech and ECAPA-TDNN model from <a href='https://github.com/speechbrain/speechbrain' target='_blank'><b>SpeechBrain</b></a> to encode and clasify speakers</h2>
            </div>
        ''')

        with gr.Row():
            gr.Markdown('''
            ### Transcribe youtube link using OpenAI Whisper
            ##### 1. Using Open AI's Whisper model to seperate audio into segments and generate transcripts.
            ##### 2. Generating speaker embeddings for each segments.
            ##### 3. Applying agglomerative clustering on the embeddings to identify the speaker for each segment.
            ''')             

        with gr.Row():
            with gr.Column():
                microphone_in.render()
                upload_in.render()
                with gr.Column():
                    gr.Markdown('''
                    ##### Here you can start the transcription process.
                    ##### Please select the source language for transcription.
                    ##### You should select a number of speakers for getting better results.
                    ''')
                selected_source_lang_audio.render()
                selected_whisper_model_audio.render()
                number_speakers_audio.render()
                transcribe_btn = gr.Button("Transcribe audio and initiate diarization")
                transcribe_btn.click(speech_to_text, 
                                    [
                                        microphone_in,
                                        upload_in,
                                        selected_source_lang_audio,
                                        selected_whisper_model_audio,
                                        number_speakers_audio
                                    ],
                                    [
                                        transcription_df_audio,
                                        system_info_audio
                                    ])

                
        with gr.Row():
            gr.Markdown('''
            ##### Here you will get transcription  output
            ##### ''')
            

        with gr.Row():
            with gr.Column():
                transcription_df_audio.render()
                system_info_audio.render()
        
        with gr.Row():
            with gr.Column():
                download_btn = gr.Button("Download transcription dataframe")
                download_btn.click(download_csv, transcription_df_audio, csv_download_audio)
                csv_download_audio.render()
            
        with gr.Row():
            gr.Markdown('''Chair of Data Science and Natural Language Processing - University of St. Gallen''')
    
    with gr.Tab("Youtube Speech to Text"):
        with gr.Row():
            gr.Markdown('''
                        <div>
                        <h1 style='text-align: center'>Youtube Speech Recognition & Speaker Diarization</h1>
                        </div>
                        ''')
            
        with gr.Row():
            gr.Markdown('''
                        ### Transcribe Youtube link
                        #### Test with the following examples:
                        ''')
            examples = gr.Examples(examples = 
                                   [
                                       "https://www.youtube.com/watch?v=vnc-Q8V4ihQ",
                                       "https://www.youtube.com/watch?v=_B60aTHCE5E",
                                       "https://www.youtube.com/watch?v=4BdKZxD-ziA",
                                       "https://www.youtube.com/watch?v=4ezBjAW26Js",
                                   ],
                                   label="Examples UNISG",
                                   inputs=[youtube_url_in])
        
        with gr.Row():
            with gr.Column():
                youtube_url_in.render()
                download_youtube_btn = gr.Button("Download Youtube video")
                download_youtube_btn.click(get_youtube, [youtube_url_in], [video_in])
                print(video_in)
        
        with gr.Row():
            with gr.Column():
                video_in.render()
                with gr.Column():
                    gr.Markdown('''
                                #### Start the transcription process.
                                #### To initiate, please select the source language for transcription.
                                #### For better performance select the number of speakers.
                                ''')
                selected_source_lang_yt.render()
                selected_whisper_model_yt.render()
                number_speakers_yt.render()
                transcribe_btn = gr.Button("Transcribe audio and initiate diarization")
                transcribe_btn.click(yt_to_text, 
                                    [
                                        video_in,
                                        selected_source_lang_yt,
                                        selected_whisper_model_yt,
                                        number_speakers_yt
                                    ],
                                    [
                                        transcription_df_yt,
                                        system_info_yt
                                    ])
        
        with gr.Row():
            gr.Markdown('''
                        #### Here you will get transcription  output
                        #### ''')
        
        with gr.Row():
            with gr.Column():
                transcription_df_yt.render()
                system_info_yt.render()

        with gr.Row():
            with gr.Column():
                download_btn = gr.Button("Download transcription dataframe")
                download_btn.click(download_csv, transcription_df_audio, csv_download_yt)
                csv_download_yt.render()
        
        with gr.Row():
            gr.Markdown('''Chair of Data Science and Natural Language Processing - University of St. Gallen''')


def main():
    interface.launch()


if __name__ == "__main__":
    main()