speaker_embed_pyannote.py

import os
import glob
import pickle
from pyannote.audio import Model, Inference
from pyannote.metrics.binary_classification import det_curve
import numpy as np
from tqdm import tqdm

# 미리 학습된 화자 임베딩 모델을 불러옵니다.
model = Model.from_pretrained("pyannote/embedding")
# window="whole" 옵션은 파일 전체에 대한 단일 임베딩을 추출합니다.
inference = Inference(model, window="whole")

def compute_embedding(audio_file):
    """
    주어진 음성 파일로부터 화자 임베딩을 산출하는 함수입니다.
    """
    embedding = inference(audio_file)
    return embedding.flatten()  # 필요에 따라 flatten 처리

def extract_features_from_audio_directory(audio_directory, output_pickle):
    """
    주어진 오디오 디렉토리에서 모든 .wav 파일을 처리하여 임베딩을 계산한 후,
    딕셔너리 형태로 저장하고 pickle 파일로 저장하는 함수입니다.

    매개변수:
      audio_directory : 오디오 파일들이 위치한 폴더(문자열)
      output_pickle   : 저장할 pickle 파일의 이름 및 경로(문자열)
    """
    embeddings_dict = {}
    
    # 디렉토리 내의 모든 .wav 파일들을 가져옵니다.
    audio_files = list(glob.glob(os.path.join(audio_directory, "*.wav")))
    
    if not audio_files:
        print("오디오 폴더 내에 .wav 파일이 없습니다.")
        return

    for audio_file in tqdm(audio_files):
        try:
            embedding = compute_embedding(audio_file)
            filename = os.path.basename(audio_file).replace(".wav", "")
            embeddings_dict[filename] = embedding
        except Exception as e:
            print(f"{audio_file} 처리 중 오류 발생: {e}")

    with open(output_pickle, "wb") as f:
        pickle.dump(embeddings_dict, f)
    print(f"모든 피쳐가 '{output_pickle}'에 저장되었습니다.")

if __name__ == "__main__":
    audio_directory = "audio"  # 오디오 파일들이 위치한 폴더
    output_pickle = "speaker_embed_pyannote.pkl"  # 저장될 pickle 파일 이름
    extract_features_from_audio_directory(audio_directory, output_pickle)