app.py

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)

    segments = transcript['segments']
    embedding_model = PretrainedSpeakerEmbedding("speechbrain/spkrec-ecapa-voxceleb", device='cpu')
    embeddings = np.zeros(shape=(len(segments), 192))

    with contextlib.closing(wave.open(output_file, 'r')) as f:
        frames = f.getnframes()
        rate = f.getframerate()
        duration = frames / float(rate)

    for i, segment in enumerate(segments):
        embeddings[i] = segment_embedding(output_file, duration, segment, embedding_model)

    embeddings = np.nan_to_num(embeddings)
    return embeddings

import numpy as np
from sklearn.cluster import AgglomerativeClustering
from typing import List, Tuple

def clustering_embeddings(speaker_count: int, embeddings: np.ndarray) -> AgglomerativeClustering:
    """
    埋め込みデータをクラスタリングして、クラスタリングオブジェクトを返します。

    Parameters
    ----------
    embeddings: np.ndarray
        分散表現（埋め込み）のリスト。

    Returns
    -------
    AgglomerativeClustering
        クラスタリングオブジェクト。
    """
    clustering = AgglomerativeClustering(speaker_count).fit(embeddings)
    return clustering

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"]))

    output = ""
    for speaker, _, text in labeled_segments:
        output += f"話者{speaker + 1}: 「{text}」\n"
    return output

from sklearn.cluster import KMeans
from sklearn.metrics.pairwise import pairwise_distances
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

from scipy.spatial.distance import cosine

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
        最も近い参照話者の名前
    """
    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

    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
        発話者ごとに整形されたセグメントの文字列。
    """
    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="音声ファイルをアップロードすると、各話者の名前がアサインされた文字起こしが作成されます。参考: https://huggingface.co/spaces/vumichien/Whisper_speaker_diarization"
)

app.launch(debug=True)