yutakobayashi/diet-members-voice-embeddings

pip install pandas numpy wave ast scipy pyannote.audio

import pandas as pd
import numpy as np
import contextlib
import wave
import ast
from typing import List, Tuple
from scipy.spatial.distance import cosine
from pyannote.audio import Audio
from pyannote.core import Segment
from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding

class SpeakerRecognizer:
    def __init__(self, threshold: float = 0.5, embedding_model_path: str = "speechbrain/spkrec-ecapa-voxceleb"):
        self.threshold = threshold
        self.embedding_model = PretrainedSpeakerEmbedding(embedding_model_path, device="cpu")
        self.audio = Audio()

    def recognize(self, target_embedding: np.ndarray, known_speaker_embeddings: List[Tuple[str, np.ndarray]]) -> str:
        distances = [cosine(target_embedding, emb) for _, emb in known_speaker_embeddings]
        min_distance, recognized_speaker = min((val, spk) for (spk, _), val in zip(known_speaker_embeddings, distances))
        return recognized_speaker if min_distance <= self.threshold else "不明な話者"

    def embed(self, audio_path: str, duration: float, segment: pd.Series) -> np.ndarray:
        segment = segment.copy()
        segment["start"] /= 1000
        segment["end"] /= 1000
        clip = Segment(segment["start"], min(duration, segment["end"]))
        waveform, _ = self.audio.crop(audio_path, clip)
        return self.embedding_model(waveform[None])

    @staticmethod
    def load_dataframes(transcript_fp: str, reference_fp: str) -> Tuple[pd.DataFrame, List[Tuple[str, np.ndarray]]]:
        transcript = pd.read_csv(transcript_fp)
        reference = pd.read_csv(reference_fp)
        known_speaker_embeddings = [(row["name"], np.array(ast.literal_eval(row["vec"]))) for _, row in reference.iterrows()]
        return transcript, known_speaker_embeddings

    @staticmethod
    def compute_duration(audio_path: str) -> float:
        with contextlib.closing(wave.open(str(audio_path), "r")) as f:
            frames = f.getnframes()
            rate = f.getframerate()
            return frames / float(rate)

    def process(self, transcript_fp: str, wav_fp: str, reference_fp: str):
        transcript, known_speaker_embeddings = self.load_dataframes(transcript_fp, reference_fp)
        duration = self.compute_duration(wav_fp)
        embeddings = np.vstack([self.embed(wav_fp, duration, segment) for _, segment in transcript.iterrows()])
        embeddings = np.nan_to_num(embeddings)

        labeled_segments = [
            (segment.start, segment.end, self.recognize(embedding, known_speaker_embeddings), segment.text)
            for embedding, segment in zip(embeddings, transcript.itertuples())
        ]

        output_by_segment = pd.DataFrame(labeled_segments, columns=["start", "end", "speaker", "text"])
        output_by_segment.to_csv("artifact.csv", index=False)

if __name__ == "__main__":
    recognizer = SpeakerRecognizer()
    recognizer.process("./transcript.csv","./audio.wav","../data/members.csv")