|
|
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 |
