|
|
import os |
|
|
import sys |
|
|
import torch |
|
|
|
|
|
import numpy as np |
|
|
import torch.nn.functional as F |
|
|
|
|
|
from functools import cached_property |
|
|
from torch.nn.utils.rnn import pad_sequence |
|
|
|
|
|
sys.path.append(os.getcwd()) |
|
|
|
|
|
from main.library.speaker_diarization.speechbrain import EncoderClassifier |
|
|
|
|
|
class BaseInference: |
|
|
pass |
|
|
|
|
|
class SpeechBrainPretrainedSpeakerEmbedding(BaseInference): |
|
|
def __init__(self, embedding = "assets/models/speaker_diarization/models/speechbrain", device = None): |
|
|
super().__init__() |
|
|
|
|
|
self.embedding = embedding |
|
|
self.device = device or torch.device("cpu") |
|
|
self.classifier_ = EncoderClassifier.from_hparams(source=self.embedding, run_opts={"device": self.device}) |
|
|
|
|
|
def to(self, device): |
|
|
if not isinstance(device, torch.device): raise TypeError |
|
|
|
|
|
self.classifier_ = EncoderClassifier.from_hparams(source=self.embedding, run_opts={"device": device}) |
|
|
self.device = device |
|
|
return self |
|
|
|
|
|
@cached_property |
|
|
def sample_rate(self): |
|
|
return self.classifier_.audio_normalizer.sample_rate |
|
|
|
|
|
@cached_property |
|
|
def dimension(self): |
|
|
*_, dimension = self.classifier_.encode_batch(torch.rand(1, 16000).to(self.device)).shape |
|
|
return dimension |
|
|
|
|
|
@cached_property |
|
|
def metric(self): |
|
|
return "cosine" |
|
|
|
|
|
@cached_property |
|
|
def min_num_samples(self): |
|
|
with torch.inference_mode(): |
|
|
lower, upper = 2, round(0.5 * self.sample_rate) |
|
|
middle = (lower + upper) // 2 |
|
|
|
|
|
while lower + 1 < upper: |
|
|
try: |
|
|
_ = self.classifier_.encode_batch(torch.randn(1, middle).to(self.device)) |
|
|
upper = middle |
|
|
except RuntimeError: |
|
|
lower = middle |
|
|
|
|
|
middle = (lower + upper) // 2 |
|
|
|
|
|
return upper |
|
|
|
|
|
def __call__(self, waveforms, masks = None): |
|
|
batch_size, num_channels, num_samples = waveforms.shape |
|
|
assert num_channels == 1 |
|
|
|
|
|
waveforms = waveforms.squeeze(dim=1) |
|
|
|
|
|
if masks is None: |
|
|
signals = waveforms.squeeze(dim=1) |
|
|
wav_lens = signals.shape[1] * torch.ones(batch_size) |
|
|
else: |
|
|
batch_size_masks, _ = masks.shape |
|
|
assert batch_size == batch_size_masks |
|
|
|
|
|
imasks = F.interpolate(masks.unsqueeze(dim=1), size=num_samples, mode="nearest").squeeze(dim=1) > 0.5 |
|
|
signals = pad_sequence([waveform[imask].contiguous() for waveform, imask in zip(waveforms, imasks)], batch_first=True) |
|
|
wav_lens = imasks.sum(dim=1) |
|
|
|
|
|
max_len = wav_lens.max() |
|
|
if max_len < self.min_num_samples: return np.nan * np.zeros((batch_size, self.dimension)) |
|
|
|
|
|
too_short = wav_lens < self.min_num_samples |
|
|
wav_lens = wav_lens / max_len |
|
|
wav_lens[too_short] = 1.0 |
|
|
|
|
|
embeddings = (self.classifier_.encode_batch(signals, wav_lens=wav_lens).squeeze(dim=1).cpu().numpy()) |
|
|
embeddings[too_short.cpu().numpy()] = np.nan |
|
|
|
|
|
return embeddings |
