speaker_test.py

import sys
import torch
import torchaudio
from pathlib import Path

# hubconf.py에 정의된 모델 로드 함수를 사용합니다.
from hubconf import load_custom
# compute_eer.py의 classify_speaker_cosine 함수를 import합니다.
from compute_eer import classify_speaker_cosine

def load_audio(audio_path, target_sample_rate=16000):
    """
    오디오 파일을 로드하고 모노 변환 및 필요시 리샘플링을 수행합니다.
    """
    waveform, sample_rate = torchaudio.load(audio_path)
    # 모노로 변환 (스테레오인 경우 채널 평균)
    if waveform.shape[0] > 1:
        waveform = torch.mean(waveform, dim=0, keepdim=True)
    # 리샘플링 (샘플레이트가 다를 경우)
    if sample_rate != target_sample_rate:
        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_sample_rate)
        waveform = resampler(waveform)
    return waveform

def main():
    if len(sys.argv) != 3:
        print("사용법: python speaker_test.py audio1.wav audio2.wav")
        sys.exit(1)
    
    wav_path1 = sys.argv[1]
    wav_path2 = sys.argv[2]
    
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print(f"사용할 디바이스: {device}")
    
    # S-vb2+vox2+cnc-ft_mix.pt 모델 로드 (hubconf.py의 load_custom 함수 사용)
    model = load_custom(model_name='S', train_type='ft_mix', dataset='vb2+vox2+cnc').to(device)
    model.eval()
    
    # 오디오 파일 로드 및 리샘플링 수행
    waveform1 = load_audio(wav_path1)
    waveform2 = load_audio(wav_path2)
    
    # 배치 차원 추가 및 디바이스로 이동
    waveform1 = waveform1.unsqueeze(0).to(device)
    waveform2 = waveform2.unsqueeze(0).to(device)
    
    # 모델 추론 (각각의 임베딩 계산)
    with torch.no_grad():
        emb1 = model(waveform1)
        emb2 = model(waveform2)
    
    # CPU로 이동 후 numpy 배열로 변환
    emb1_np = emb1.cpu().numpy()
    emb2_np = emb2.cpu().numpy()
    
    # threshold 값 0.856 사용하여 같은 화자인지 판별 (compute_eer.py의 classify_speaker_cosine 활용)
    is_same = classify_speaker_cosine(emb1_np, emb2_np, 0.856)
    
    if is_same:
        print("두 오디오 파일은 같은 화자입니다.")
    else:
        print("두 오디오 파일은 다른 화자입니다.")

if __name__ == "__main__":
    main()