Update app.py

app.py

@@ -1,288 +1,295 @@
+# import whisper
+from faster_whisper import WhisperModel
+import datetime
+import subprocess
+import gradio as gr
+from pathlib import Path
+import pandas as pd
+import re
+import time
 import os
-import wave
 import numpy as np
-import contextlib
-from pydub import AudioSegment
-from pyannote.core import Segment
-from pyannote.audio import Audio
-from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding
-import torch
-from typing import Dict, List, Tuple
-
-
-def convert_to_wav(input_file: str, output_file: str = "output_file.wav") -> str:
-    """
-    音声ファイルをWAV形式に変換します。
-
-    Parameters
-    ----------
-    input_file: str
-        変換する音声ファイルのパス
-    output_file: str
-        変換後のWAVファイルの出力先パス（デフォルトは"output_file.wav"）
-    Returns
-    -------
-    str
-        変換後のWAVファイルのパス
-    """
-    file_format = os.path.splitext(input_file)[1][1:]
-    audio = AudioSegment.from_file(input_file, format=file_format)
-    audio.export(output_file, format="wav")
-    return output_file
-
-
-def segment_embedding(
-        file_name: str,
-        duration: float,
-        segment,
-        embedding_model: PretrainedSpeakerEmbedding
-) -> np.ndarray:
-    """
-    音声ファイルから指定されたセグメントの埋め込みを計算します。
-
-    Parameters
-    ----------
-    file_name: str
-        音声ファイルのパス
-    duration: float
-        音声ファイルの継続時間
-    segment: whisperのtranscribeのsegment
-    embedding_model: PretrainedSpeakerEmbedding
-        埋め込みモデル
-    Returns
-    -------
-    np.ndarray
-        計算された埋め込みベクトル
-    """
-    audio = Audio()
-    start = segment["start"]
-    end = min(duration, segment["end"])
-    clip = Segment(start, end)
-    waveform, sample_rate = audio.crop(file_name, clip)
-    return embedding_model(waveform[None])
-
-
-def reference_audio_embedding(
-        file_name: str
-) -> np.ndarray:
-    """
-    参考音声の埋め込みを出力します。
-
-    Parameters
-    ----------
-    file_name: str
-        音声ファイルのパス
-    Returns
-    -------
-    np.ndarray
-        計算された埋め込みベクトル
-    """
-    audio = Audio()
-    waveform, sample_rate = audio(file_name)
-    embedding_model = embedding_model = PretrainedSpeakerEmbedding("speechbrain/spkrec-ecapa-voxceleb", device='cpu')
-    return embedding_model(waveform[None])[0]
-
-
-def generate_speaker_embeddings(
-        meeting_file_path: str,
-        transcript
-) -> np.ndarray:
-    """
-    音声ファイルから話者の埋め込みを計算します。
-
-    Parameters
-    ----------
-    meeting_file_path: str
-        音声ファイルのパス
-    transcript: Whisper API の transcribe メソッドの出力結果
-    Returns
-    -------
-    np.ndarray
-        計算された話者の埋め込み群
-    """
-    output_file = convert_to_wav(meeting_file_path)
+from sklearn.cluster import AgglomerativeClustering
+from sklearn.metrics import silhouette_score
 
-    segments = transcript['segments']
-    embedding_model = PretrainedSpeakerEmbedding("speechbrain/spkrec-ecapa-voxceleb", device='cpu')
-    embeddings = np.zeros(shape=(len(segments), 192))
+from pytube import YouTube
+import yt_dlp
+import torch
+import pyannote.audio
+from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding
+from pyannote.audio import Audio
+from pyannote.core import Segment
 
-    with contextlib.closing(wave.open(output_file, 'r')) as f:
-        frames = f.getnframes()
-        rate = f.getframerate()
-        duration = frames / float(rate)
+from gpuinfo import GPUInfo
 
-    for i, segment in enumerate(segments):
-        embeddings[i] = segment_embedding(output_file, duration, segment, embedding_model)
+import wave
+import contextlib
+from transformers import pipeline
+import psutil
+
+whisper_models = ["tiny", "base", "small", "medium", "large-v1", "large-v2"]
+source_languages = {
+    "en": "English",
+    "ja": "Japanese",
+}
+
+source_language_list = [key[0] for key in source_languages.items()]
+
+MODEL_NAME = "vumichien/whisper-medium-jp"
+lang = "ja"
+
+device = 0 if torch.cuda.is_available() else "cpu"
+pipe = pipeline(
+    task="automatic-speech-recognition",
+    model=MODEL_NAME,
+    chunk_length_s=30,
+    device=device,
+)
+os.makedirs('output', exist_ok=True)
+pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language=lang, task="transcribe")
 
-    embeddings = np.nan_to_num(embeddings)
-    return embeddings
+embedding_model = PretrainedSpeakerEmbedding(
+    "speechbrain/spkrec-ecapa-voxceleb",
+    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"))
 
 
-import numpy as np
-from sklearn.cluster import AgglomerativeClustering
-from typing import List, Tuple
+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"
 
-def clustering_embeddings(speaker_count: int, embeddings: np.ndarray) -> AgglomerativeClustering:
-    """
-    埋め込みデータをクラスタリングして、クラスタリングオブジェクトを返します。
-    Parameters
-    ----------
-    embeddings: np.ndarray
-        分散表現（埋め込み）のリスト。
-    Returns
-    -------
-    AgglomerativeClustering
-        クラスタリングオブジェクト。
-    """
-    clustering = AgglomerativeClustering(speaker_count).fit(embeddings)
-    return clustering
+    file = microphone if microphone is not None else file_upload
 
+    text = pipe(file)["text"]
 
-def format_speaker_output_by_segment(clustering: AgglomerativeClustering, transcript: dict) -> str:
-    """
-    クラスタリングの結果をもとに、各発話者ごとにセグメントを整形して出力します
-    Parameters
-    ----------
-    clustering: AgglomerativeClustering
-        クラスタリングオブジェクト。
-    transcript: dict
-        Whisper API の transcribe メソッドの出力結果
-    Returns
-    -------
-    str
-        発話者ごとに整形されたセグメントの文字列
-    """
-    labeled_segments = []
-    for label, segment in zip(clustering.labels_, transcript["segments"]):
-        labeled_segments.append((label, segment["start"], segment["text"]))
+    return warn_output + text
 
-    output = ""
-    for speaker, _, text in labeled_segments:
-        output += f"話者{speaker + 1}: 「{text}」\n"
-    return output
 
+def convert_time(secs):
+    return datetime.timedelta(seconds=round(secs))
 
-from sklearn.cluster import KMeans
-from sklearn.metrics.pairwise import pairwise_distances
 
+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()
 
-def clustering_embeddings2(speaker_count: int, embeddings: np.ndarray) -> KMeans:
-    """
-    埋め込みデータをクラスタリングして、クラスタリングオブジェクトを返します。
-    Parameters
-    ----------
-    embeddings: np.ndarray
-        分散表現（埋め込み）のリスト。
-    Returns
-    -------
-    KMeans
-        クラスタリングオブジェクト。
-    """
-    # コサイン類似度行列を計算
-    cosine_distances = pairwise_distances(embeddings, metric='cosine')
-    clustering = KMeans(n_clusters=speaker_count).fit(cosine_distances)
-    return clustering
+    ydl_opts = {
+        'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best',
+    }
 
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        info = ydl.extract_info(video_url, download=False)
+        abs_video_path = ydl.prepare_filename(info)
+        ydl.process_info(info)
 
-from scipy.spatial.distance import cosine
+    print("Success download video")
+    print(abs_video_path)
+    return abs_video_path
 
 
-def closest_reference_speaker(embedding: np.ndarray, references: List[Tuple[str, np.ndarray]]) -> str:
-    """
-    与えられた埋め込みに最も近い参照話者を返します。
-    Parameters
-    ----------
-    embedding: np.ndarray
-        話者の埋め込み
-    references: List[Tuple[str, np.ndarray]]
-        参照話者の名前と埋め込みのリスト
-    Returns
-    -------
-    str
-        最も近い参照話者の名前
+def speech_to_text(video_file_path, selected_source_lang, whisper_model, num_speakers):
     """
-    min_distance = float('inf')
-    closest_speaker = None
-    for name, reference_embedding in references:
-        distance = cosine(embedding, reference_embedding)
-        if distance < min_distance:
-            min_distance = distance
-            closest_speaker = name
+    # 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.
 
-    return closest_speaker
-
-
-def format_speaker_output_by_segment2(embeddings: np.ndarray, transcript: dict,
-                                      reference_embeddings: List[Tuple[str, np.ndarray]]) -> str:
-    """
-    各発話者の埋め込みに基づいて、セグメントを整形して出力します。
-    Parameters
-    ----------
-    embeddings: np.ndarray
-        話者の埋め込みのリスト
-    transcript: dict
-        Whisper API の transcribe メソッドの出力結果
-    reference_embeddings: List[Tuple[str, np.ndarray]]
-        参照話者の名前と埋め込みのリスト
-    Returns
-    -------
-    str
-        発話者ごとに整形されたセグメントの文字列。
+    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
     """
-    labeled_segments = []
-    for embedding, segment in zip(embeddings, transcript["segments"]):
-        speaker_name = closest_reference_speaker(embedding, reference_embeddings)
-        labeled_segments.append((speaker_name, segment["start"], segment["text"]))
-
-    output = ""
-    for speaker, _, text in labeled_segments:
-        output += f"{speaker}: 「{text}」\n"
-    return output
-
-
-import gradio as gr
-import openai
-
-
-def create_transcription_with_speaker(openai_key, main_audio, reference_audio_1, reference1_name,
-                                      reference_audio_2, reference2_name, speaker_count=2):
-    openai.api_key = openai_key
-    # 文字起こし
-    transcript = openai.Audio.transcribe("whisper-1", open(main_audio, "rb"), response_format="verbose_json")
-    # 各発話をembeddingsに変換
-    embeddings = generate_speaker_embeddings(main_audio, transcript)
-    # 各発話のembeddingsをクラスタリング
-    clustering = clustering_embeddings(speaker_count, embeddings)
-    # クラスタリングで作られた仮のラベルで各セグメントに名前付け
-    output_by_segment1 = format_speaker_output_by_segment(clustering, transcript)
-    reference1 = reference_audio_embedding(reference_audio_1)
-    reference2 = reference_audio_embedding(reference_audio_2)
-    reference_embeddings = [(reference1_name, reference1), (reference2_name, reference2)]
-    output_by_segment2 = format_speaker_output_by_segment2(embeddings, transcript, reference_embeddings)
-    return output_by_segment1, output_by_segment2
-
-
-inputs = [
-    gr.Textbox(lines=1, label="openai_key", type="password"),
-    gr.Audio(type="filepath", label="メイン音声ファイル"),
-    gr.Audio(type="filepath", label="話者 (1) 参考音声ファイル"),
-    gr.Textbox(lines=1, label="話者 (1) の名前"),
-    gr.Audio(type="filepath", label="話者 (2) 参考音声ファイル"),
-    gr.Textbox(lines=1, label="話者 (2) の名前")
-]
-
-outputs = [
-    gr.Textbox(label="話者クラスタリング文字起こし"),
-    gr.Textbox(label="話者アサイン文字起こし"),
-]
-
-app = gr.Interface(
-    fn=create_transcription_with_speaker,
-    inputs=inputs,
-    outputs=outputs,
-    title="話者アサイン機能付き書き起こしアプリ",
-    description="音声ファイルをアップロードすると、話者分離した文字起こしが作成されます。"
-)
 
-app.launch(debug=True)
+    # model = whisper.load_model(whisper_model)
+    # model = WhisperModel(whisper_model, device="cuda", compute_type="int8_float16")
+    model = WhisperModel(whisper_model, compute_type="int8")
+    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 enging 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)
+        segments_raw, info = model.transcribe(audio_file, **transcribe_options)
+
+        # Convert back to original openai format
+        segments = []
+        i = 0
+        for segment_chunk in segments_raw:
+            chunk = {}
+            chunk["start"] = segment_chunk.start
+            chunk["end"] = segment_chunk.end
+            chunk["text"] = segment_chunk.text
+            segments.append(chunk)
+            i += 1
+        print("transcribe audio done with fast 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}')
+
+        if num_speakers == 0:
+            # Find the best number of speakers
+            score_num_speakers = {}
+
+            for num_speakers in range(2, 10 + 1):
+                clustering = AgglomerativeClustering(num_speakers).fit(embeddings)
+                score = silhouette_score(embeddings, clustering.labels_, metric='euclidean')
+                score_num_speakers[num_speakers] = score
+            best_num_speaker = max(score_num_speakers, key=lambda x: score_num_speakers[x])
+            print(f"The best number of speakers: {best_num_speaker} with {score_num_speakers[best_num_speaker]} score")
+        else:
+            best_num_speaker = num_speakers
+
+        # Assign speaker label
+        clustering = AgglomerativeClustering(best_num_speaker).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 = ''
+        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.*
+        """
+        save_path = "output/transcript_result.csv"
+        df_results = pd.DataFrame(objects)
+        df_results.to_csv(save_path)
+        return df_results, system_info, save_path
+
+    except Exception as e:
+        raise RuntimeError("Error Running inference with local model", e)
+
+
+# ---- Gradio Layout -----
+# Inspiration from https://huggingface.co/spaces/RASMUS/Whisper-youtube-crosslingual-subtitles
+video_in = gr.Video(label="Video file", mirror_webcam=False)
+youtube_url_in = gr.Textbox(label="Youtube url", lines=1, interactive=True)
+df_init = pd.DataFrame(columns=['Start', 'End', 'Speaker', 'Text'])
+memory = psutil.virtual_memory()
+selected_source_lang = gr.Dropdown(choices=source_language_list, type="value", value="ja",
+                                   label="Spoken language in video", interactive=True)
+selected_whisper_model = gr.Dropdown(choices=whisper_models, type="value", value="base", label="Selected Whisper model",
+                                     interactive=True)
+number_speakers = gr.Number(precision=0, value=0,
+                            label="Input number of speakers for better results. If value=0, model will automatic find the best number of speakers",
+                            interactive=True)
+system_info = gr.Markdown(
+    f"*Memory: {memory.total / (1024 * 1024 * 1024):.2f}GB, used: {memory.percent}%, available: {memory.available / (1024 * 1024 * 1024):.2f}GB*")
+download_transcript = gr.File(label="Download transcript")
+transcription_df = gr.DataFrame(value=df_init, label="Transcription dataframe", row_count=(0, "dynamic"), max_rows=10,
+                                wrap=True, overflow_row_behaviour='paginate')
+title = "Whisper speaker diarization"
+demo = gr.Blocks(title=title)
+demo.encrypt = False
+
+with demo:
+    with gr.Row():
+        gr.Markdown('''
+            ### You can test by following examples:
+            ''')
+        examples = gr.Examples(examples=
+                           ["https://www.youtube.com/watch?v=j7BfEzAFuYc&t=32s",
+                            "https://www.youtube.com/watch?v=-UX0X45sYe4",
+                            "https://www.youtube.com/watch?v=7minSgqi-Gw"],
+                           label="Examples", 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('''
+                ##### Here you can start the transcription process.
+                ##### Please select the source language for transcription.
+                ##### You can select a range of assumed numbers of speakers.
+                ''')
+            selected_source_lang.render()
+            selected_whisper_model.render()
+            number_speakers.render()
+            transcribe_btn = gr.Button("Transcribe audio and diarization")
+            transcribe_btn.click(speech_to_text,
+                                 [video_in, selected_source_lang, selected_whisper_model, number_speakers],
+                                 [transcription_df, system_info, download_transcript]
+                                 )
+
+    with gr.Row():
+        gr.Markdown('''
+        ##### Here you will get transcription  output
+        ##### ''')
+
+    with gr.Row():
+        with gr.Column():
+            download_transcript.render()
+            transcription_df.render()
+            # system_info.render()
+
+
+demo.launch(debug=True)