C++ version of pyannote audio speaker diarizaiton pipeline
import numpy as np
from einops import rearrange
from scipy.cluster.hierarchy import fcluster, linkage
from scipy.optimize import linear_sum_assignment
from scipy.spatial.distance import cdist, pdist


class BaseClustering(Pipeline):
    def __init__(
        self,
        metric: str = "cosine",
        max_num_embeddings: int = 1000,
        constrained_assignment: bool = False,
    ):  

        super().__init__()
        self.metric = metric
        self.max_num_embeddings = max_num_embeddings
        self.constrained_assignment = constrained_assignment

    def set_num_clusters(
        self,
        num_embeddings: int,
        num_clusters: int = None,
        min_clusters: int = None,
        max_clusters: int = None,
    ):  

        min_clusters = num_clusters or min_clusters or 1
        min_clusters = max(1, min(num_embeddings, min_clusters))
        max_clusters = num_clusters or max_clusters or num_embeddings
        max_clusters = max(1, min(num_embeddings, max_clusters))

        if min_clusters > max_clusters:
            raise ValueError(
                f"min_clusters must be smaller than (or equal to) max_clusters "
                f"(here: min_clusters={min_clusters:g} and max_clusters={max_clusters:g})."
            )   

        if min_clusters == max_clusters:
            num_clusters = min_clusters

        return num_clusters, min_clusters, max_clusters


    def filter_embeddings(
        self,
        embeddings: np.ndarray,
        segmentations: SlidingWindowFeature = None,
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Filter NaN embeddings and downsample embeddings

        Parameters
        ----------
        embeddings : (num_chunks, num_speakers, dimension) array
            Sequence of embeddings.
        segmentations : (num_chunks, num_frames, num_speakers) array
            Binary segmentations.

        Returns
        -------
        filtered_embeddings : (num_embeddings, dimension) array
        chunk_idx : (num_embeddings, ) array
        speaker_idx : (num_embeddings, ) array
        """

        chunk_idx, speaker_idx = np.where(~np.any(np.isnan(embeddings), axis=2))

        # sample max_num_embeddings embeddings
        num_embeddings = len(chunk_idx)
        if num_embeddings > self.max_num_embeddings:
            indices = list(range(num_embeddings))
            random.shuffle(indices)
            indices = sorted(indices[: self.max_num_embeddings])
            chunk_idx = chunk_idx[indices]
            speaker_idx = speaker_idx[indices]

        return embeddings[chunk_idx, speaker_idx], chunk_idx, speaker_idx


    def constrained_argmax(self, soft_clusters: np.ndarray) -> np.ndarray:

        soft_clusters = np.nan_to_num(soft_clusters, nan=np.nanmin(soft_clusters))
        num_chunks, num_speakers, num_clusters = soft_clusters.shape
        # num_chunks, num_speakers, num_clusters

        hard_clusters = -2 * np.ones((num_chunks, num_speakers), dtype=np.int8)

        for c, cost in enumerate(soft_clusters):
            speakers, clusters = linear_sum_assignment(cost, maximize=True)
            for s, k in zip(speakers, clusters):
                hard_clusters[c, s] = k

        return hard_clusters


    def assign_embeddings(
        self,
        embeddings: np.ndarray,
        train_chunk_idx: np.ndarray,
        train_speaker_idx: np.ndarray,
        train_clusters: np.ndarray,
        constrained: bool = False,
    ):
        """Assign embeddings to the closest centroid

        Cluster centroids are computed as the average of the train embeddings
        previously assigned to them.

        Parameters
        ----------
        embeddings : (num_chunks, num_speakers, dimension)-shaped array
            Complete set of embeddings.
        train_chunk_idx : (num_embeddings,)-shaped array
        train_speaker_idx : (num_embeddings,)-shaped array
            Indices of subset of embeddings used for "training".
        train_clusters : (num_embedding,)-shaped array
            Clusters of the above subset
        constrained : bool, optional
            Use constrained_argmax, instead of (default) argmax.

        Returns
        -------
        soft_clusters : (num_chunks, num_speakers, num_clusters)-shaped array
        hard_clusters : (num_chunks, num_speakers)-shaped array
        """

        # TODO: option to add a new (dummy) cluster in case num_clusters < max(frame_speaker_count)

        num_clusters = np.max(train_clusters) + 1
        num_chunks, num_speakers, dimension = embeddings.shape

        train_embeddings = embeddings[train_chunk_idx, train_speaker_idx]

        centroids = np.vstack(
            [
                np.mean(train_embeddings[train_clusters == k], axis=0)
                for k in range(num_clusters)
            ]
        )

        # compute distance between embeddings and clusters
        e2k_distance = rearrange(
            cdist(
                rearrange(embeddings, "c s d -> (c s) d"),
                centroids,
                metric=self.metric,
            ),
            "(c s) k -> c s k",
            c=num_chunks,
            s=num_speakers,
        )
        soft_clusters = 2 - e2k_distance

        # assign each embedding to the cluster with the most similar centroid
        if constrained:
            hard_clusters = self.constrained_argmax(soft_clusters)
        else:
            hard_clusters = np.argmax(soft_clusters, axis=2)

        # TODO: add a flag to revert argmax for trainign subset
        # hard_clusters[train_chunk_idx, train_speaker_idx] = train_clusters

        return hard_clusters, soft_clusters


    def __call__(
        self,
        embeddings: np.ndarray,
        segmentations: SlidingWindowFeature = None,
        num_clusters: int = None,
        min_clusters: int = None,
        max_clusters: int = None,
        **kwargs,
    ) -> np.ndarray:
        """Apply clustering

        Parameters
        ----------
        embeddings : (num_chunks, num_speakers, dimension) array
            Sequence of embeddings.
        segmentations : (num_chunks, num_frames, num_speakers) array
            Binary segmentations.
        num_clusters : int, optional
            Number of clusters, when known. Default behavior is to use
            internal threshold hyper-parameter to decide on the number
            of clusters.
        min_clusters : int, optional
            Minimum number of clusters. Has no effect when `num_clusters` is provided.
        max_clusters : int, optional
            Maximum number of clusters. Has no effect when `num_clusters` is provided.

        Returns
        -------
        hard_clusters : (num_chunks, num_speakers) array
            Hard cluster assignment (hard_clusters[c, s] = k means that sth speaker
            of cth chunk is assigned to kth cluster)
        soft_clusters : (num_chunks, num_speakers, num_clusters) array
            Soft cluster assignment (the higher soft_clusters[c, s, k], the most likely
            the sth speaker of cth chunk belongs to kth cluster)
        """

        train_embeddings, train_chunk_idx, train_speaker_idx = self.filter_embeddings(
            embeddings,
            segmentations=segmentations,
        )

        num_embeddings, _ = train_embeddings.shape
        num_clusters, min_clusters, max_clusters = self.set_num_clusters(
            num_embeddings,
            num_clusters=num_clusters,
            min_clusters=min_clusters,
            max_clusters=max_clusters,
        )

        if max_clusters < 2:
            # do NOT apply clustering when min_clusters = max_clusters = 1
            num_chunks, num_speakers, _ = embeddings.shape
            hard_clusters = np.zeros((num_chunks, num_speakers), dtype=np.int8)
            soft_clusters = np.ones((num_chunks, num_speakers, 1))
            return hard_clusters, soft_clusters

        train_clusters = self.cluster(
            train_embeddings,
            min_clusters,
            max_clusters,
            num_clusters=num_clusters,
        )

        hard_clusters, soft_clusters = self.assign_embeddings(
            embeddings,
            train_chunk_idx,
            train_speaker_idx,
            train_clusters,
            constrained=self.constrained_assignment,
        )

        return hard_clusters, soft_clusters


class AgglomerativeClustering(BaseClustering):
    """Agglomerative clustering

    Parameters
    ----------
    metric : {"cosine", "euclidean", ...}, optional
        Distance metric to use. Defaults to "cosine".

    Hyper-parameters
    ----------------
    method : {"average", "centroid", "complete", "median", "single", "ward"}
        Linkage method.
    threshold : float in range [0.0, 2.0]
        Clustering threshold.
    """

    def __init__(
        self,
        metric: str = "cosine",
        max_num_embeddings: int = np.inf,
        constrained_assignment: bool = False,
    ):

        super().__init__(
            metric=metric,
            max_num_embeddings=max_num_embeddings,
            constrained_assignment=constrained_assignment,
        )

        self.threshold = Uniform(0.0, 2.0)  # assume unit-normalized embeddings
        self.method = Categorical(
            ["average", "centroid", "complete", "median", "single", "ward", "weighted"]
        )

        # minimum cluster size
        self.min_cluster_size = Integer(1, 20)

    def cluster(
        self,
        embeddings: np.ndarray,
        min_clusters: int,
        max_clusters: int,
        num_clusters: int = None,
    ):
       """

        Parameters
        ----------
        embeddings : (num_embeddings, dimension) array
            Embeddings
        min_clusters : int
            Minimum number of clusters
        max_clusters : int
            Maximum number of clusters
        num_clusters : int, optional
            Actual number of clusters. Default behavior is to estimate it based
            on values provided for `min_clusters`,  `max_clusters`, and `threshold`.

        Returns
        -------
        clusters : (num_embeddings, ) array
            0-indexed cluster indices.
        """

        num_embeddings, _ = embeddings.shape

        # heuristic to reduce self.min_cluster_size when num_embeddings is very small
        # (0.1 value is kind of arbitrary, though)
        min_cluster_size = min(
            self.min_cluster_size, max(1, round(0.1 * num_embeddings))
        )

        # linkage function will complain when there is just one embedding to cluster
        if num_embeddings == 1:
            return np.zeros((1,), dtype=np.uint8)

        # centroid, median, and Ward method only support "euclidean" metric
        # therefore we unit-normalize embeddings to somehow make them "euclidean"
        if self.metric == "cosine" and self.method in ["centroid", "median", "ward"]:
            with np.errstate(divide="ignore", invalid="ignore"):
                embeddings /= np.linalg.norm(embeddings, axis=-1, keepdims=True)
            dendrogram: np.ndarray = linkage(
                embeddings, method=self.method, metric="euclidean"
            )

        # other methods work just fine with any metric
        else:
            dendrogram: np.ndarray = linkage(
                embeddings, method=self.method, metric=self.metric
            )

       # apply the predefined threshold
        clusters = fcluster(dendrogram, self.threshold, criterion="distance") - 1

        # split clusters into two categories based on their number of items:
        # large clusters vs. small clusters
        cluster_unique, cluster_counts = np.unique(
            clusters,
            return_counts=True,
        )
        large_clusters = cluster_unique[cluster_counts >= min_cluster_size]
        num_large_clusters = len(large_clusters)

        # force num_clusters to min_clusters in case the actual number is too small
        if num_large_clusters < min_clusters:
            num_clusters = min_clusters

        # force num_clusters to max_clusters in case the actual number is too large
        elif num_large_clusters > max_clusters:
            num_clusters = max_clusters

        if num_clusters is not None:

            # switch stopping criterion from "inter-cluster distance" stopping to "iteration index"
            _dendrogram = np.copy(dendrogram)
            _dendrogram[:, 2] = np.arange(num_embeddings - 1)

            best_iteration = num_embeddings - 1
            best_num_large_clusters = 1

            # traverse the dendrogram by going further and further away
            # from the "optimal" threshold

            for iteration in np.argsort(np.abs(dendrogram[:, 2] - self.threshold)):

                # only consider iterations that might have resulted
                # in changing the number of (large) clusters
                new_cluster_size = _dendrogram[iteration, 3]
                if new_cluster_size < min_cluster_size:
                    continue

                # estimate number of large clusters at considered iteration
                clusters = fcluster(_dendrogram, iteration, criterion="distance") - 1
                cluster_unique, cluster_counts = np.unique(clusters, return_counts=True)
                large_clusters = cluster_unique[cluster_counts >= min_cluster_size]
                num_large_clusters = len(large_clusters)
               # keep track of iteration that leads to the number of large clusters
                # as close as possible to the target number of clusters.
                if abs(num_large_clusters - num_clusters) < abs(
                    best_num_large_clusters - num_clusters
                ):
                    best_iteration = iteration
                    best_num_large_clusters = num_large_clusters

                # stop traversing the dendrogram as soon as we found a good candidate
                if num_large_clusters == num_clusters:
                    break

            # re-apply best iteration in case we did not find a perfect candidate
            if best_num_large_clusters != num_clusters:
                clusters = (
                    fcluster(_dendrogram, best_iteration, criterion="distance") - 1
                )
                cluster_unique, cluster_counts = np.unique(clusters, return_counts=True)
                large_clusters = cluster_unique[cluster_counts >= min_cluster_size]
                num_large_clusters = len(large_clusters)
                print(
                    f"Found only {num_large_clusters} clusters. Using a smaller value than {min_cluster_size} for `min_cluster_size` might help."
                )

        if num_large_clusters == 0:
            clusters[:] = 0
            return clusters

        small_clusters = cluster_unique[cluster_counts < min_cluster_size]
        if len(small_clusters) == 0:
            return clusters

        # re-assign each small cluster to the most similar large cluster based on their respective centroids
        large_centroids = np.vstack(
            [
                np.mean(embeddings[clusters == large_k], axis=0)
                for large_k in large_clusters
            ]
        )
        small_centroids = np.vstack(
            [
                np.mean(embeddings[clusters == small_k], axis=0)
                for small_k in small_clusters
            ]
        )
        centroids_cdist = cdist(large_centroids, small_centroids, metric=self.metric)
        for small_k, large_k in enumerate(np.argmin(centroids_cdist, axis=0)):
            clusters[clusters == small_clusters[small_k]] = large_clusters[large_k]

        # re-number clusters from 0 to num_large_clusters
        _, clusters = np.unique(clusters, return_inverse=True)
        return clusters



clustering/Clustering.py
leohuang2013/pyannote-audio_speaker-diarization_cpp
import torch
import torch.nn as nn
import torch.nn.functional as F

''' Res2Conv1d + BatchNorm1d + ReLU
'''
class Res2Conv1dReluBn(nn.Module):
    '''
    in_channels == out_channels == channels
    '''
    def __init__(self, channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False, scale=4):
        super().__init__()
        assert channels % scale == 0, "{} % {} != 0".format(channels, scale)
        self.scale = scale
        self.width = channels // scale
        self.nums = scale if scale == 1 else scale - 1

        self.convs = []
        self.bns = []
        for i in range(self.nums):
            self.convs.append(nn.Conv1d(self.width, self.width, kernel_size, stride, padding, dilation, bias=bias))
            self.bns.append(nn.BatchNorm1d(self.width))
        self.convs = nn.ModuleList(self.convs)
        self.bns = nn.ModuleList(self.bns)

    def forward(self, x):
        out = []
        spx = torch.split(x, self.width, 1)
        for i in range(self.nums):
            if i == 0:
                sp = spx[i]
            else:
                sp = sp + spx[i]
            # Order: conv -> relu -> bn
            sp = self.convs[i](sp)
            sp = self.bns[i](F.relu(sp))
            out.append(sp)
        if self.scale != 1:
            out.append(spx[self.nums])
        out = torch.cat(out, dim=1)
        return out



''' Conv1d + BatchNorm1d + ReLU
'''
class Conv1dReluBn(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False):
        super().__init__()
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)
        self.bn = nn.BatchNorm1d(out_channels)

    def forward(self, x):
        return self.bn(F.relu(self.conv(x)))



''' The SE connection of 1D case.
'''
class SE_Connect(nn.Module):
    def __init__(self, channels, s=2):
        super().__init__()
        assert channels % s == 0, "{} % {} != 0".format(channels, s)
        self.linear1 = nn.Linear(channels, channels // s)
        self.linear2 = nn.Linear(channels // s, channels)

    def forward(self, x):
        out = x.mean(dim=2)
        out = F.relu(self.linear1(out))
        out = torch.sigmoid(self.linear2(out))
        out = x * out.unsqueeze(2)
        return out



''' SE-Res2Block.
    Note: residual connection is implemented in the ECAPA_TDNN model, not here.
'''
def SE_Res2Block(channels, kernel_size, stride, padding, dilation, scale):
    return nn.Sequential(
        Conv1dReluBn(channels, channels, kernel_size=1, stride=1, padding=0),
        Res2Conv1dReluBn(channels, kernel_size, stride, padding, dilation, scale=scale),
        Conv1dReluBn(channels, channels, kernel_size=1, stride=1, padding=0),
        SE_Connect(channels)
    )



''' Attentive weighted mean and standard deviation pooling.
'''
class AttentiveStatsPool(nn.Module):
    def __init__(self, in_dim, bottleneck_dim):
        super().__init__()
        # Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs.
        self.linear1 = nn.Conv1d(in_dim, bottleneck_dim, kernel_size=1) # equals W and b in the paper
        self.linear2 = nn.Conv1d(bottleneck_dim, in_dim, kernel_size=1) # equals V and k in the paper

    def forward(self, x):
        # DON'T use ReLU here! In experiments, I find ReLU hard to converge.
        alpha = torch.tanh(self.linear1(x))
        alpha = torch.softmax(self.linear2(alpha), dim=2)
        mean = torch.sum(alpha * x, dim=2)
        residuals = torch.sum(alpha * x ** 2, dim=2) - mean ** 2
        std = torch.sqrt(residuals.clamp(min=1e-9))
        return torch.cat([mean, std], dim=1)



''' Implementation of
    "ECAPA-TDNN: Emphasized Channel Attention, Propagation and Aggregation in TDNN Based Speaker Verification".

    Note that we DON'T concatenate the last frame-wise layer with non-weighted mean and standard deviation, 
    because it brings little improvment but significantly increases model parameters. 
    As a result, this implementation basically equals the A.2 of Table 2 in the paper.
'''
class ECAPA_TDNN(nn.Module):
    def __init__(self, in_channels=80, channels=512, embd_dim=192):
        super().__init__()
        self.layer1 = Conv1dReluBn(in_channels, channels, kernel_size=5, padding=2)
        self.layer2 = SE_Res2Block(channels, kernel_size=3, stride=1, padding=2, dilation=2, scale=8)
        self.layer3 = SE_Res2Block(channels, kernel_size=3, stride=1, padding=3, dilation=3, scale=8)
        self.layer4 = SE_Res2Block(channels, kernel_size=3, stride=1, padding=4, dilation=4, scale=8)

        cat_channels = channels * 3
        self.conv = nn.Conv1d(cat_channels, 1536, kernel_size=1)
        self.pooling = AttentiveStatsPool(1536, 128)
        self.bn1 = nn.BatchNorm1d(3072)
        self.linear = nn.Linear(3072, embd_dim)
        self.bn2 = nn.BatchNorm1d(embd_dim)

    def forward(self, x):
        x = x.transpose(1, 2)
        out1 = self.layer1(x)
        out2 = self.layer2(out1) + out1
        out3 = self.layer3(out1 + out2) + out1 + out2
        out4 = self.layer4(out1 + out2 + out3) + out1 + out2 + out3

        out = torch.cat([out2, out3, out4], dim=1)
        out = F.relu(self.conv(out))
        out = self.bn1(self.pooling(out))
        out = self.bn2(self.linear(out))
        return out
      


if __name__ == '__main__':
    # Input size: batch_size * seq_len * feat_dim
    x = torch.zeros(2, 200, 80)
    model = ECAPA_TDNN(in_channels=80, channels=512, embd_dim=192)
    out = model(x)
    print(model)
    print(out.shape)    # should be [2, 192]

    # Export the model
    symbolic_names = {0: "batch_size", 1: "max_seq_len"}
    torch.onnx.export(model,               # model being run
                      x,                         # model input (or a tuple for multiple inputs)
                      "ecapa-tdnn.onnx",   # where to save the model (can be a file or file-like object)
                      export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=17,          # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      verbose=False,
                      input_names = ['signal'],   # the model's input names
                      output_names = ['embeddings'], # the model's output names
                      dynamic_axes={'signal' : symbolic_names,    # variable length axes
                                    })
    # Check exported model see if there is error
    import onnx
    onnx_model = onnx.load("ecapa-tdnn.onnx")
    onnx.checker.check_model(onnx_model)

    # Compare the output of the original model and the ONNX-converted model to ensure their equivalence.
    import onnxruntime
    rep = onnx.shape_inference.infer_shapes(onnx_model)
    #print( rep )
    providers = ['CUDAExecutionProvider']
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), providers=providers)
    ort_inputs = {ort_session.get_inputs()[0].name: x.numpy()}
    ort_outs = ort_session.run(None, ort_inputs)
    # Check output is same( or close enough )
    #np.testing.assert_allclose(embeddings.detach().numpy(), ort_outs[0], rtol=1e-03, atol=1e-05)
    #print("Original Output:", embeddings)
    print("Onnx model Output:", ort_outs[0].shape)


embeddings/ECAPA-TDNN.py
leohuang2013/pyannote-audio_speaker-diarization_cpp
from pyannote.audio import Pipeline
from hf_token import *
import datetime
import time
import sys
import wave
import contextlib

import rich


def diarize( wavFile, min_speakers = None, max_speakers = None ):

    try:
        with contextlib.closing(wave.open(wavFile,'r')) as f:
            frames = f.getnframes()
            rate = f.getframerate()
            duration = frames / float(rate)
        print( f'Diarize wav file: {wavFile} with duration: {duration}' )
    except Exception as e:
        print( f'Failed to open wav file: {e}' )
        return None
    
    # get the start datetime
    st = datetime.datetime.now()

    print( "load model" )
    pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization",
                                        use_auth_token=hf_auth_token)
    #pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization",
    #                                    cache_dir="/Users/leo/.cache/torch/pyannote")

    print( "apply pipeline" )
    # apply pretrained pipeline
    if min_speakers == None or max_speakers == None or \
         min_speakers == 0 or max_speakers == 0:
        '''
        If do not know the number of speaker 
        '''
        diarization = pipeline( wavFile )
    else:
        print( f"speakers info: min={min_speakers}, max={max_speakers}" )
        if min_speakers == max_speakers or min_speakers > max_speakers:
            '''
            If know exactly the number of speaker
            '''
            diarization = pipeline("audio.wav", num_speakers=min_speakers)
        elif min_speakers > max_speakers: # wrong config
            diarization = pipeline( wavFile )
        else:
            '''
            If know min and max of the number of speaker 
            '''
            diarization = pipeline("audio.wav", min_speakers=min_speakers, max_speakers=max_speakers)

    print('----------------------------------------------------')
    result = []
    # print the result
    for turn, _, speaker in diarization.itertracks(yield_label=True):
        print(f"start={turn.start:.1f}s stop={turn.end:.1f}s speaker_{speaker}")
        result.append({ "start": round(turn.start, 2), "end": round(turn.end, 2), "speaker": f"{speaker}" })
    print('----------------------------------------------------')

    # get the end datetime
    et = datetime.datetime.now()

    # get execution time
    elapsed_time = et - st
    print('Time cost:', elapsed_time, 'seconds')

    return result

def main():
    if len( sys.argv ) < 2 :
        print( 'Wav file needed' )
        return

    wavFile = sys.argv[1]
    diarize( wavFile )

if __name__ == '__main__':
    main()

embeddings/SpeakerDiarizer.py
leohuang2013/pyannote-audio_speaker-diarization_cpp
import torchaudio
from speechbrain.pretrained import EncoderClassifier
import torch
import numpy as np

model = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
        savedir="pretrained")

signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav')
#signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/multi-speaker_7-speakers_SkillsFuture.wav')
print( signal )
print( signal.shape )
#exit( 0 ) 

wavs = signal
if len(wavs.shape) == 1:
    wavs = wavs.unsqueeze(0)
wav_lens = torch.ones(wavs.shape[0], device=model.device)
wavs = wavs.float()
feats = model.mods.compute_features(wavs)
feats = model.mods.mean_var_norm(feats, wav_lens)

# Print output shape
#embeddings = model.encode_batch(signal)
embeddings = model.encode_batch(signal, normalize=False)
#print( embeddings )
#print( embeddings.shape )
#exit( 0 ) 
#prediction = model.classify_batch(signal)
#print(prediction)
#exit(0)


def export():
    # Create dummy input
    # Batch of waveforms [batch, time, channels] or [batch, time]
    # speechbrain/pretrained/interfaces.py:903
    #x = torch.randn( 1, 1920000 )
    x = torch.randn( 1, 200, 80 )
    print( x )
    print( x.shape )
    #exit( 0 ) 

    pretrained_dict = model.state_dict()
    #print( pretrained_dict )
    #exit(0)

    # Show all unconvertable ops, output would be like,
    # {'aten::view_as_real', 'aten::chunk', 'aten::stft'}
    torch_script_graph, unconvertible_ops = torch.onnx.utils.unconvertible_ops(
        model, x, opset_version=17
    )
    if( len( unconvertible_ops ) > 0 ):
        print( '------------------------There are some unconvertable operations-----------' )
        print(set(unconvertible_ops))
        exit( 0 )
    print( '---- all operations convertable ----' )

    # Export the model
    #symbolic_names = {0: "batch_size", 1: "max_seq_len"}
    symbolic_names = {0: "batch_size", 1: "max_seq_len", 2: "dim"}
    torch.onnx.export(model,               # model being run
                      x,                         # model input (or a tuple for multiple inputs)
                      "embeddings.onnx",   # where to save the model (can be a file or file-like object)
                      export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=17,          # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      verbose=False,
                      input_names = ['signal'],   # the model's input names
                      output_names = ['embeddings'], # the model's output names
                      dynamic_axes={'signal' : symbolic_names,    # variable length axes
                                    })

def verify():
    # Check exported model see if there is error
    import onnx
    onnx_model = onnx.load("embeddings.onnx")
    onnx.checker.check_model(onnx_model)

    # Compare the output of the original model and the ONNX-converted model to ensure their equivalence.
    import onnxruntime
    rep = onnx.shape_inference.infer_shapes(onnx_model)
    #print(rep)
    providers = ['CPUExecutionProvider']
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), providers=providers)
    print( ort_session.get_inputs()[0])
    #ort_inputs = {ort_session.get_inputs()[0].name: signal.numpy()}
    ort_inputs = {ort_session.get_inputs()[0].name: feats.numpy()}
    ort_outs = ort_session.run(None, ort_inputs)
    # Check output is same( or close enough )
    np.testing.assert_allclose(embeddings.detach().numpy(), ort_outs[0], rtol=1e-03, atol=1e-05)
    print("Original Output:", embeddings)
    print("Onnx model Output:", ort_outs[0])


#export()
verify()

embeddings/export.py
leohuang2013/pyannote-audio_speaker-diarization_cpp
import torchaudio
from speechbrain.pretrained import EncoderClassifier
import torch
import numpy as np

model = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
        savedir="pretrained")

signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav')
#signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/multi-speaker_7-speakers_SkillsFuture.wav')
print( signal )
print( signal.shape )
#exit( 0 ) 

wavs = signal
if len(wavs.shape) == 1:
    wavs = wavs.unsqueeze(0)
wav_lens = torch.ones(wavs.shape[0], device=model.device)
wavs = wavs.float()
feats = model.mods.compute_features(wavs)
#feats = model.mods.mean_var_norm(feats, wav_lens)

# Print output shape
#embeddings = model.encode_batch(signal)
embeddings = model.encode_batch(signal, normalize=False)
#print( embeddings )
#print( embeddings.shape )
#exit( 0 ) 
#prediction = model.classify_batch(signal)
#print(prediction)
#exit(0)

outModel = 'embeddings1.onnx'

def export():
    # Create dummy input
    # Batch of waveforms [batch, time, channels] or [batch, time]
    # speechbrain/pretrained/interfaces.py:903
    x = torch.randn( 1, 1920000 )
    #x = torch.randn( 1, 200, 80 )
    print( x )
    print( x.shape )
    #exit( 0 ) 

    pretrained_dict = model.state_dict()
    #print( pretrained_dict )
    #exit(0)

    # Show all unconvertable ops, output would be like,
    # {'aten::view_as_real', 'aten::chunk', 'aten::stft'}
    torch_script_graph, unconvertible_ops = torch.onnx.utils.unconvertible_ops(
        model, x, opset_version=17
    )
    if( len( unconvertible_ops ) > 0 ):
        print( '------------------------There are some unconvertable operations-----------' )
        print(set(unconvertible_ops))
        exit( 0 )
    print( '---- all operations convertable ----' )

    # Export the model
    symbolic_names = {0: "batch_size", 1: "max_seq_len"}
    #symbolic_names = {0: "batch_size", 1: "max_seq_len", 2: "dim"}
    torch.onnx.export(model,               # model being run
                      x,                         # model input (or a tuple for multiple inputs)
                      outModel,   # where to save the model (can be a file or file-like object)
                      export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=17,          # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      verbose=False,
                      input_names = ['signal'],   # the model's input names
                      output_names = ['embeddings'], # the model's output names
                      dynamic_axes={'signal' : symbolic_names,    # variable length axes
                                    })

def verify():
    # Check exported model see if there is error
    import onnx
    onnx_model = onnx.load( outModel )
    onnx.checker.check_model(onnx_model)

    # Compare the output of the original model and the ONNX-converted model to ensure their equivalence.
    import onnxruntime
    rep = onnx.shape_inference.infer_shapes(onnx_model)
    #print(rep)
    providers = ['CPUExecutionProvider']
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), providers=providers)
    print( ort_session.get_inputs()[0])
    ort_inputs = {ort_session.get_inputs()[0].name: signal.numpy()}
    #ort_inputs = {ort_session.get_inputs()[0].name: feats.numpy()}
    ort_outs = ort_session.run(None, ort_inputs)
    # Check output is same( or close enough )
    #np.testing.assert_allclose(embeddings.detach().numpy(), ort_outs[0], rtol=1e-03, atol=1e-05)
    #np.testing.assert_allclose(feats.detach().numpy(), ort_outs[0], rtol=1e-03, atol=1e-05,
    #        verbose=True)
    import sys
    np.set_printoptions(threshold=sys.maxsize)
    print("Original Output:")
    #print(feats.detach().numpy())
    a1 = feats.detach().numpy().flatten()
    print("Onnx model Output:")
    #print(ort_outs[0])
    a2 = ort_outs[0].flatten()
    if len( a1 ) != len( a2 ):
        print( f"Size diff: {len(a1)}, {len(a2)}" )
    for i in range( len( a1 )):
        if abs( a1[i] - a2[i] ) > 1e-03:
            print( f"Diff: {a1[i]}, {a2[i]}" )
            return


export()
verify()

embeddings/export1.py
leohuang2013/pyannote-audio_speaker-diarization_cpp
import torchaudio
import onnx
import onnxruntime
from myembedding import MyEmbedding
import torch
import numpy as np
import torch.onnx.verification
from speechbrain.pretrained import EncoderClassifier


outModel = 'embeddings2.onnx'
opsetVer = 17


def simple():
    model = MyEmbedding.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")

    signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav')
    #signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav')
    #signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/multi-speaker_7-speakers_SkillsFuture.wav')
    print( signal )
    print( signal.shape )
    #exit( 0 ) 

    wavs = signal
    if len(wavs.shape) == 1:
        wavs = wavs.unsqueeze(0)
    wav_lens = torch.ones(wavs.shape[0], device=model.device)
    wavs, wav_lens = wavs.to(model.device), wav_lens.to(model.device)
    wavs = wavs.float()
    feats = model.mods.compute_features(wavs)
    print( feats )
    #f = open( '/tmp/feats1.txt', 'w' )
    #df = feats.numpy().flatten()
    #for a in df:
    #    f.write( str( a ) + ',' )
    #f.close()
    feats = model.mods.mean_var_norm(feats, wav_lens)
    #exit( 0 ) 

    # Print output shape
    embeddings = model.encode_batch(signal, normalize=False)
    #print( embeddings )
    #print( embeddings.shape )
    #exit( 0 ) 
    #prediction = model.classify_batch(signal)
    #print(prediction)
    #exit(0)

def export( ):
    # Create dummy input
    # Batch of waveforms [batch, time, channels] or [batch, time]
    # speechbrain/pretrained/interfaces.py:903
    batch = 32
    #batch = 1
    max_sequence = 80000
    #x = torch.FloatTensor(batch, max_sequence).uniform_(-32768, 32767)
    x = torch.FloatTensor(batch, max_sequence).uniform_(-1, 1)
    print( x.shape )
    #exit( 0 ) 

    model = MyEmbedding.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")
    model.eval()
    pretrained_dict = model.state_dict()
    #print( pretrained_dict )
    #exit(0)

    # TODO: test if eval() will solve precision problem
    # another potential cause is, one is on cpu, another is on cuda
    #model.eval()

    # Show all unconvertable ops, output would be like,
    # {'aten::view_as_real', 'aten::chunk', 'aten::stft'}
    torch_script_graph, unconvertible_ops = torch.onnx.utils.unconvertible_ops(
        model, 
        x, 
        opset_version=opsetVer
    )
    if( len( unconvertible_ops ) > 0 ):
        print( '------------------------There are some unconvertable operations-----------' )
        print(set(unconvertible_ops))
        exit( 0 )
    print( '\n---- all operations convertable ----\n' )

    #print( '\n---- Check model mismatch ----' )
    #torch.onnx.verification.find_mismatch( model,
    #        x,
    #        opset_version=opsetVer,
    #        do_constant_folding=True,
    #        verbose=True
    #        )
    #print( '---- Mismatch check is done ----\n' )

    signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav')
    print( '\n---- start export ----' )
    # Export the model
    '''
    There will be big difference between using signal and x for result of mean_var_norm
    '''
    #model_script = torch.jit.script( model )
    wav_lens = torch.Tensor([1.0])
    wav_lens1 = [1.0] * batch
    wav_lens1 = torch.Tensor( wav_lens1 )
    wav_lens1 = wav_lens1.to(torch.float32)
    torch.onnx.export( model,               # model being run
                      #(signal, wav_lens),                          # model input (or a tuple for multiple inputs)
                      (x, wav_lens1),                          # model input (or a tuple for multiple inputs)
                      outModel,   # where to save the model (can be a file or file-like object)
                      export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=opsetVer,          # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      verbose=False,
                      input_names = ['signal', 'wav_lens'],   # the model's input names
                      output_names = ['embeddings'], # the model's output names
                      dynamic_axes={
                          'signal': {0: "batch_size", 1: "max_seq_len"},    
                          'wav_lens': {0: "batch_size" }
                                    })
    print( f'---- model has been saved to {outModel} ----\n' )

def debug():
    '''
    Load model
    check model
    print output shape
    print output 
    print outputs of all layers -- > this is very useful for debugging, 
        - check name against graph in netron.app 
        - python -m pdb export2.py to check each step input/output
    '''
    # Check exported model see if there is error
    # add all intermediate outputs to onnx net ------ 1
    ort_session = onnxruntime.InferenceSession(outModel)
    org_outputs = [x.name for x in ort_session.get_outputs()]

    print( f'\n---- Load model: {outModel} ----' )
    onnx_model = onnx.load( outModel )
    # Check model
    onnx.checker.check_model(onnx_model, full_check=True)
    inferred = onnx.shape_inference.infer_shapes(onnx_model, check_type=True)

    # ----- 2
    for node in onnx_model.graph.node:
        for output in node.output:
            if output not in org_outputs:
                onnx_model.graph.output.extend([onnx.ValueInfoProto(name=output)])

    # Load test data
    #signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav')
    signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/one_min16k.wav')
    print( signal.shape )
    #model = MyEmbedding.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
    model = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")

    # Calc embeddings using pytorch
    wavs = signal
    if len(wavs.shape) == 1:
        wavs = wavs.unsqueeze(0)
    wav_lens = torch.ones(wavs.shape[0], device=model.device)
    wavs, wav_lens = wavs.to(model.device), wav_lens.to(model.device)
    wavs = wavs.float()
    feats = model.mods.compute_features(wavs)
    feats = model.mods.mean_var_norm(feats, wav_lens)
    d1 = signal.numpy()
    embeddings = model.encode_batch(signal, normalize=False)
    d2 = signal.numpy()
    # make sure above function call, signal not changed
    np.testing.assert_allclose(
            d1, 
            d2, 
            rtol=1e-03, 
            atol=1e-05,
            verbose = True)
    print( '\nSignal data is not changed\n' )

    # Compare the output of the original model and the ONNX-converted model to ensure their equivalence.
    providers = ['CPUExecutionProvider']
    so = onnxruntime.SessionOptions()
    so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
    #so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_BASIC
    so.log_severity_level = 1 
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), so, providers=providers)
    print( ort_session.get_inputs()[0])
    ort_inputs = {ort_session.get_inputs()[0].name: signal.numpy()}
    outputs = [x.name for x in ort_session.get_outputs()]
    #ort_inputs = {ort_session.get_inputs()[0].name: feats.numpy()}

    print( f'---- run model ----' )
    #ort_outs = ort_session.run(None, ort_inputs)
    ort_outs = ort_session.run(outputs, ort_inputs)
    onnx_embeddings = ort_outs[0]
    # ------ 3
    from collections import OrderedDict
    ort_outs = OrderedDict(zip(outputs, ort_outs))
    # Check output is same( or close enough )

    print( '---- print all layers outputs ---- ' )
    for name in outputs:
        print( f'--- out for {name} ---' )
        print( ort_outs[name] )

    print( f'---- check result ----' )
    #fn = torch.from_numpy( ort_outs[0] )
    #out = model.mods.mean_var_norm( fn, wav_lens)
    np.testing.assert_allclose(
            embeddings.detach().numpy(), 
            #feats.detach().numpy(), 
            onnx_embeddings, 
            rtol=1e-03, 
            atol=1e-05,
            verbose = True)
    #np.testing.assert_allclose(
    #        feats.detach().numpy(), 
    #        ort_outs[0], 
    #        rtol=1e-03, 
    #        atol=1e-05,
    #        verbose=True)
    print( f'---- check result is done ----' )

    import sys
    np.set_printoptions(threshold=sys.maxsize)
    print("Original Output:")
    #print(feats.detach().numpy())
    a1 = feats.detach().numpy().flatten()
    print("Onnx model Output:")
    #print(ort_outs[0])
    a2 = ort_outs[0].flatten()
    if len( a1 ) != len( a2 ):
        print( f"Size diff: {len(a1)}, {len(a2)}" )
    for i in range( len( a1 )):
        if abs( a1[i] - a2[i] ) > 1e-03:
            print( f"Diff: {a1[i]}, {a2[i]}" )
            return


def verify():
    model = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")
    #model = MyEmbedding.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
    #        savedir="pretrained")
    #model.eval()
    onnx_model = onnx.load( outModel )
    providers = ['CPUExecutionProvider']
    so = onnxruntime.SessionOptions()
    so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), so, providers=providers)

    wavs = [
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/multi-speaker_4-speakers_Jennifer_Aniston_and_Adam_Sandler_talk.wav',
            #'/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav',
            #'/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav',
            ]

    batch = 32
    for wav in wavs:
        print( f"Check wav file: {wav}" )

        # Load wav
        signal, fs = torchaudio.load( wav )
        signal = signal[:,:batch*80000].flatten().reshape( batch, -1 )

        # Calc embeddings using pytorch
        print( signal.shape )
        wav_lens = torch.Tensor([1.0] * batch)
        embeddings = model.encode_batch(signal, wav_lens, normalize=False)

        # Calc embedding using onnx model
        print( ort_session.get_inputs()[0])
        print( ort_session.get_inputs()[1])
        ort_inputs = {ort_session.get_inputs()[0].name: signal.numpy(),
                ort_session.get_inputs()[1].name: wav_lens.numpy()
                }
        ort_outs = ort_session.run(None, ort_inputs)
        onnx_embeddings = ort_outs[0]

        try:
            # Check if result is close enough
            np.testing.assert_allclose(
                    embeddings.detach().numpy(), 
                    onnx_embeddings, 
                    rtol=1e-03, 
                    atol=1e-05,
                    verbose = True)
            print( 'check passed' )
        except AssertionError as e:
            print( e )
        print( "-----------------------------" )

def test():
    onnx_model = onnx.load( outModel )
    providers = ['CPUExecutionProvider']
    so = onnxruntime.SessionOptions()
    so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), so, providers=providers)

    # Load wav
    with open( '/tmp/cpp_signals0.txt' ) as f:
        lines = f.readlines()
    signals = []
    for line in lines:
        arr = line.split(',')
        arr = arr[:-1]
        assert( len( arr ) == 80000 )
        signal = [float(a) for a in arr ]
        signals.append( signal )
    signals = torch.Tensor( signals )

    with open( '/tmp/cpp_final_wav_lens0.txt' ) as f:
        content = f.read()
    arr = content.split(',')
    arr = arr[:-1]
    assert( len( arr ) == signals.shape[0])
    wav_lens = [float(a) for a in arr ]
    # Calc embedding using onnx model
    print( ort_session.get_inputs()[0])
    print( ort_session.get_inputs()[1])
    wav_lens = torch.Tensor( wav_lens )
    ort_inputs = {ort_session.get_inputs()[0].name: signals.numpy(),
            ort_session.get_inputs()[1].name: wav_lens.numpy()
            }
    ort_outs = ort_session.run(None, ort_inputs)
    onnx_embeddings = ort_outs[0]
    import sys
    np.set_printoptions(threshold=sys.maxsize)
    print( onnx_embeddings )


export()
#debug()
verify()
#test()
#simple()



embeddings/export2.py
leohuang2013/pyannote-audio_speaker-diarization_cpp
import torchaudio
import onnx
import onnxruntime
from threeModel import FeatureModel, EMDModel, FBankModel, MySTFT, MyEmbedding0
import torch
import numpy as np
import torch.onnx.verification
from speechbrain.pretrained import EncoderClassifier
from speechbrain.processing.features import STFT, Filterbank, spectral_magnitude


opsetVer = 17


def export( outModel ):
    # Create dummy input
    # Batch of waveforms [batch, time, channels] or [batch, time]
    # speechbrain/pretrained/interfaces.py:903
    #batch = 32
    batch = 1
    max_sequence = 80000
    #x = torch.FloatTensor(batch, max_sequence).uniform_(-32768, 32767)
    x = torch.FloatTensor(batch, max_sequence).uniform_(-1, 1)
    print( x.shape )
    #exit( 0 ) 

    model = FeatureModel.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")
    model.eval()
    pretrained_dict = model.state_dict()
    #print( pretrained_dict )
    #exit(0)

    # TODO: test if eval() will solve precision problem
    # another potential cause is, one is on cpu, another is on cuda
    #model.eval()

    # Show all unconvertable ops, output would be like,
    # {'aten::view_as_real', 'aten::chunk', 'aten::stft'}
    torch_script_graph, unconvertible_ops = torch.onnx.utils.unconvertible_ops(
        model, 
        x, 
        opset_version=opsetVer
    )
    if( len( unconvertible_ops ) > 0 ):
        print( '------------------------There are some unconvertable operations-----------' )
        print(set(unconvertible_ops))
        exit( 0 )
    print( '\n---- all operations convertable ----\n' )

    #print( '\n---- Check model mismatch ----' )
    #torch.onnx.verification.find_mismatch( model,
    #        x,
    #        opset_version=opsetVer,
    #        do_constant_folding=True,
    #        verbose=True
    #        )
    #print( '---- Mismatch check is done ----\n' )

    signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav')
    print( '\n---- start export ----' )
    # Export the model
    '''
    There will be big difference between using signal and x for result of mean_var_norm
    '''
    #model_script = torch.jit.script( model )
    torch.onnx.export( model,               # model being run
                      signal,                          # model input (or a tuple for multiple inputs)
                      #x,                          # model input (or a tuple for multiple inputs)
                      outModel,   # where to save the model (can be a file or file-like object)
                      export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=opsetVer,          # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      verbose=False,
                      input_names = ['signal'],   # the model's input names
                      output_names = ['features'], # the model's output names
                      dynamic_axes={
                          'signal': {0: "batch_size", 1: "max_seq_len"},    
                                    })
    print( f'---- model has been saved to {outModel} ----\n' )

def export1( outModel ):
    # Create dummy input
    #batch = 32
    batch = 1
    max_sequence = 80
    #x = torch.FloatTensor(batch, max_sequence).uniform_(-32768, 32767)
    x = torch.FloatTensor(batch, 6000, max_sequence).uniform_(-1, 1)
    print( x.shape )
    #exit( 0 ) 

    model = EMDModel.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")
    model.eval()
    pretrained_dict = model.state_dict()
    #print( pretrained_dict )
    #exit(0)
    wav_lens = torch.Tensor([1.0])
    wav_lens1 = [1.0] * batch
    wav_lens1 = torch.Tensor( wav_lens1 )
    wav_lens1 = wav_lens1.to(torch.float32)

    # TODO: test if eval() will solve precision problem
    # another potential cause is, one is on cpu, another is on cuda
    #model.eval()

    # Show all unconvertable ops, output would be like,
    # {'aten::view_as_real', 'aten::chunk', 'aten::stft'}
    torch_script_graph, unconvertible_ops = torch.onnx.utils.unconvertible_ops(
        model, 
        (x, wav_lens1), 
        opset_version=opsetVer
    )
    if( len( unconvertible_ops ) > 0 ):
        print( '------------------------There are some unconvertable operations-----------' )
        print(set(unconvertible_ops))
        exit( 0 )
    print( '\n---- all operations convertable ----\n' )

    #print( '\n---- Check model mismatch ----' )
    #torch.onnx.verification.find_mismatch( model,
    #        x,
    #        opset_version=opsetVer,
    #        do_constant_folding=True,
    #        verbose=True
    #        )
    #print( '---- Mismatch check is done ----\n' )

    signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav')
    print( '\n---- start export ----' )
    # Export the model
    '''
    There will be big difference between using signal and x for result of mean_var_norm
    '''
    #model_script = torch.jit.script( model )
    torch.onnx.export( model,               # model being run
                      #(signal, wav_lens),                          # model input (or a tuple for multiple inputs)
                      (x, wav_lens1),                          # model input (or a tuple for multiple inputs)
                      outModel,   # where to save the model (can be a file or file-like object)
                      export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=opsetVer,          # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      verbose=False,
                      input_names = ['signal'],   # the model's input names
                      output_names = ['features'], # the model's output names
                      dynamic_axes={
                          'signal': {0: "batch_size", 1: "max_seq_len"},    
                                    })
    print( f'---- model has been saved to {outModel} ----\n' )

def export2( outModel ):
    # Create dummy input
    batch = 32
    x = torch.FloatTensor(batch, 501, 201, 2).uniform_(-1, 1)
    print( x.shape )
    #exit( 0 ) 

    model = MyEmbedding0.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")
    model.eval()
    wav_lens1 = [1.0] * batch
    wav_lens1 = torch.Tensor( wav_lens1 )
    wav_lens1 = wav_lens1.to(torch.float32)

    #print( '\n---- Check model mismatch ----' )
    #torch.onnx.verification.find_mismatch( model,
    #        x,
    #        opset_version=opsetVer,
    #        do_constant_folding=True,
    #        verbose=True
    #        )
    #print( '---- Mismatch check is done ----\n' )

    print( '\n---- start export ----' )

    # Export the model
    torch.onnx.export( model,               # model being run
                      (x, wav_lens1),                          # model input (or a tuple for multiple inputs)
                      outModel,   # where to save the model (can be a file or file-like object)
                      export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=opsetVer,          # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      verbose=False,
                      input_names = ['feats', 'wav_lens'],   # the model's input names
                      output_names = ['embedding'], # the model's output names
                      dynamic_axes={
                          'feats': {0: "B", 1: "T", 2: "N", 3: "M"},    
                          'wav_lens': {0: "B" }
                                    })
    print( f'---- model has been saved to {outModel} ----\n' )

def export_stft( outModel ):
    #batch = 32
    batch = 1
    max_sequence = 80000
    x = torch.FloatTensor(batch, max_sequence).uniform_(-1, 1)
    print( x.shape )

    model = MySTFT(sample_rate=16000, win_length=25, hop_length=10, n_fft=400) 
    #model.eval()

    #print( '\n---- Check model mismatch ----' )
    #torch.onnx.verification.find_mismatch( model,
    #        x,
    #        opset_version=opsetVer,
    #        do_constant_folding=True,
    #        verbose=True
    #        )
    #print( '---- Mismatch check is done ----\n' )

    print( '\n---- start export ----' )
    torch.onnx.export( model,               # model being run
                      x,                          # model input (or a tuple for multiple inputs)
                      outModel,   # where to save the model (can be a file or file-like object)
                      #export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=opsetVer,          # the ONNX version to export the model to
                      #do_constant_folding=True,  # whether to execute constant folding for optimization
                      verbose=False,
                      input_names = ['signal'],   # the model's input names
                      output_names = ['features'], # the model's output names
                      dynamic_axes={
                          'signal': {0: "batch_size", 1: "max_seq_len"},    
                                    })
    print( f'---- model has been saved to {outModel} ----\n' )

def export_fbank( outModel ):
    #batch = 32
    batch = 1
    x = torch.FloatTensor(batch, 501, 201, 2).uniform_(-1, 1)
    print( x.shape )

    model = FBankModel() 

    #print( '\n---- Check model mismatch ----' )
    #torch.onnx.verification.find_mismatch( model,
    #        x,
    #        opset_version=opsetVer,
    #        do_constant_folding=True,
    #        verbose=True
    #        )
    #print( '---- Mismatch check is done ----\n' )

    print( '\n---- start export ----' )
    torch.onnx.export( model,               # model being run
                      x,                          # model input (or a tuple for multiple inputs)
                      outModel,   # where to save the model (can be a file or file-like object)
                      export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=opsetVer,          # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      verbose=False,
                      input_names = ['signal'],   # the model's input names
                      output_names = ['features'], # the model's output names
                      dynamic_axes={
                          'signal': {0: "B", 1: "T", 2: "N", 3: "M"},    
                                    })
    print( f'---- model has been saved to {outModel} ----\n' )

def verify_fbank( outModel ):
    model = STFT(sample_rate=16000, win_length=25, hop_length=10, n_fft=400)

    onnx_model = onnx.load( outModel )
    providers = ['CPUExecutionProvider']
    so = onnxruntime.SessionOptions()
    so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), so, providers=providers)

    wavs = [
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching1.wav',
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav',
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav',
            ]

    for wav in wavs:
        print( f"Check wav file: {wav}" )

        # Load wav
        signal, fs = torchaudio.load( wav )
        signal = signal[:,:80000]

        # Calc embeddings using pytorch
        print( signal.shape )
        compute_STFT = STFT(sample_rate=16000, win_length=25, hop_length=10, n_fft=400)
        feats_stft = compute_STFT(signal)
        features = spectral_magnitude(feats_stft)
        compute_fbanks = Filterbank(n_mels=80)
        feats = compute_fbanks(features)

        # Calc embedding using onnx model
        print( ort_session.get_inputs()[0])
        ort_inputs = {ort_session.get_inputs()[0].name: feats_stft.numpy(),
                }
        ort_outs = ort_session.run(None, ort_inputs)
        onnx_feats = ort_outs[0]

        try:
            # Check if result is close enough
            np.testing.assert_allclose(
                    feats.detach().numpy(), 
                    onnx_feats, 
                    rtol=1e-03, 
                    atol=1e-05,
                    verbose = True)
            print( 'check passed' )
        except AssertionError as e:
            print( e )
        print( "-----------------------------" )

def debug( outModel ):
    '''
    Load model
    check model
    print output shape
    print output 
    print outputs of all layers -- > this is very useful for debugging, 
        - check name against graph in netron.app 
        - python -m pdb export2.py to check each step input/output
    '''
    # Check exported model see if there is error
    # add all intermediate outputs to onnx net ------ 1
    ort_session = onnxruntime.InferenceSession(outModel)
    org_outputs = [x.name for x in ort_session.get_outputs()]

    print( f'\n---- Load model: {outModel} ----' )
    onnx_model = onnx.load( outModel )
    # Check model
    onnx.checker.check_model(onnx_model, full_check=True)
    inferred = onnx.shape_inference.infer_shapes(onnx_model, check_type=True)

    # ----- 2
    for node in onnx_model.graph.node:
        for output in node.output:
            if output not in org_outputs:
                onnx_model.graph.output.extend([onnx.ValueInfoProto(name=output)])

    # Load test data
    #signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav')
    signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav')
    print( signal.shape )
    model = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")

    # Calc embeddings using pytorch
    wavs = signal
    if len(wavs.shape) == 1:
        wavs = wavs.unsqueeze(0)
    wav_lens = torch.ones(wavs.shape[0], device=model.device)
    wavs, wav_lens = wavs.to(model.device), wav_lens.to(model.device)
    wavs = wavs.float()
    feats = model.mods.compute_features(wavs)
    #feats = model.mods.mean_var_norm(feats, wav_lens)
    #embeddings = model.encode_batch(signal, normalize=False)

    # Compare the output of the original model and the ONNX-converted model to ensure their equivalence.
    providers = ['CPUExecutionProvider']
    so = onnxruntime.SessionOptions()
    so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
    #so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_BASIC
    so.log_severity_level = 1 
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), so, providers=providers)
    print( ort_session.get_inputs()[0])
    ort_inputs = {ort_session.get_inputs()[0].name: signal.numpy()}
    outputs = [x.name for x in ort_session.get_outputs()]
    #ort_inputs = {ort_session.get_inputs()[0].name: feats.numpy()}

    print( f'---- run model ----' )
    #ort_outs = ort_session.run(None, ort_inputs)
    ort_outs = ort_session.run(outputs, ort_inputs)
    onnx_feats = ort_outs[0]
    # ------ 3
    from collections import OrderedDict
    ort_outs = OrderedDict(zip(outputs, ort_outs))
    # Check output is same( or close enough )

    print( '---- print all layers outputs ---- ' )
    for name in outputs:
        print( f'--- out for {name} ---' )
        print( ort_outs[name] )

    print( f'---- check result ----' )
    #fn = torch.from_numpy( ort_outs[0] )
    #out = model.mods.mean_var_norm( fn, wav_lens)
    np.testing.assert_allclose(
            feats.detach().numpy(), 
            #feats.detach().numpy(), 
            onnx_feats, 
            rtol=1e-03, 
            atol=1e-05,
            verbose = True)
    #np.testing.assert_allclose(
    #        feats.detach().numpy(), 
    #        ort_outs[0], 
    #        rtol=1e-03, 
    #        atol=1e-05,
    #        verbose=True)
    print( f'---- check result is done ----' )

    import sys
    np.set_printoptions(threshold=sys.maxsize)
    print("Original Output:")
    #print(feats.detach().numpy())
    a1 = feats.detach().numpy().flatten()
    print("Onnx model Output:")
    #print(ort_outs[0])
    a2 = ort_outs[0].flatten()
    if len( a1 ) != len( a2 ):
        print( f"Size diff: {len(a1)}, {len(a2)}" )
    for i in range( len( a1 )):
        if abs( a1[i] - a2[i] ) > 1e-03:
            print( f"Diff: {a1[i]}, {a2[i]}" )
            return


def verify( outModel ):
    #model = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
    #        savedir="pretrained")
    model = FeatureModel.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")
    model.eval()
    onnx_model = onnx.load( outModel )
    providers = ['CPUExecutionProvider']
    so = onnxruntime.SessionOptions()
    so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), so, providers=providers)

    wavs = [
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching1.wav',
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav',
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav',
            ]

    visualisation = {}
    def hook_fn(m, i, o):
        print( f"hook for: {m}" )
        visualisation[m] = o

    def get_all_layers(net):
        for name, layer in net._modules.items():
            #If it is a sequential, don't register a hook on it
            # but recursively register hook on all it's module children
            if isinstance(layer, torch.nn.Sequential):
                get_all_layers(layer)
            else:
                # it's a non sequential. Register a hook
                layer.register_forward_hook(hook_fn)

    get_all_layers(model.mods.compute_features)

    for wav in wavs:
        print( f"Check wav file: {wav}" )

        # Load wav
        signal, fs = torchaudio.load( wav )
        signal = signal[:,:80000]

        # Calc embeddings using pytorch
        print( signal.shape )
        wav_lens = torch.Tensor([1.0])
        feats = model.encode_batch(signal)

        print( visualisation.keys()  )

        # Calc embedding using onnx model
        print( ort_session.get_inputs()[0])
        ort_inputs = {ort_session.get_inputs()[0].name: signal.numpy(),
                }
        ort_outs = ort_session.run(None, ort_inputs)
        onnx_feats = ort_outs[0]

        try:
            # Check if result is close enough
            np.testing.assert_allclose(
                    feats.detach().numpy(), 
                    onnx_feats, 
                    rtol=1e-03, 
                    atol=1e-05,
                    verbose = True)
            print( 'check passed' )
        except AssertionError as e:
            print( e )
        print( "-----------------------------" )

def genFeats():
    model = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")
    model.eval()
    wavs = [
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching1.wav',
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav',
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav',
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/multi-speaker_1min.wav',
            ]

    n = 0
    for wav in wavs:
        # Load wav
        signal, fs = torchaudio.load( wav )
        signal = signal[:,:80000]

        # Calc embeddings using pytorch
        print( signal.shape )
        feats = model.mods.compute_features(signal)
        emd = feats.detach().numpy()
        emd = emd.flatten()
        # [1, 6001, 80]
        np.savetxt(f'feats{n}.txt', emd, fmt='%f')
        n += 1

def verify1( outModel ):
    #model = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
    #        savedir="pretrained")
    model = EMDModel.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")
    model.eval()
    onnx_model = onnx.load( outModel )
    providers = ['CPUExecutionProvider']
    so = onnxruntime.SessionOptions()
    so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), so, providers=providers)

    featsN = ['feats0.txt', 'feats1.txt', 'feats2.txt', 'feats3.txt']


    batch = 1
    wav_lens1 = [1.0] * batch
    wav_lens1 = torch.Tensor( wav_lens1 )
    wav_lens1 = wav_lens1.to(torch.float32)
    for f in featsN:
        print( "\n-----------------------------" )
        feats = np.loadtxt( f, dtype=float)
        feats = feats.reshape( 1, -1, 80 )
        feats = torch.Tensor( feats )

        # Calc embeddings using pytorch
        print( feats.shape )
        wav_lens = torch.Tensor([1.0])
        embeddings = model.encode_batch( feats )

        # Calc embedding using onnx model
        print( ort_session.get_inputs()[0])
        print( ort_session.get_inputs()[1])
        ort_inputs = {ort_session.get_inputs()[0].name: feats.numpy(),
                ort_session.get_inputs()[1].name: wav_lens1.numpy()
                }
        ort_outs = ort_session.run(None, ort_inputs)
        onnx_embeddings = ort_outs[0]

        try:
            # Check if result is close enough
            np.testing.assert_allclose(
                    embeddings.detach().numpy(), 
                    onnx_embeddings, 
                    rtol=1e-03, 
                    atol=1e-05,
                    verbose = True)
            print( f'check passed for {f}' )
        except AssertionError as e:
            print( e )
        print( "-----------------------------\n" )

def verify2( outModel ):
    model_origin = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")
    model_origin.eval()
    model = MyEmbedding0.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")
    model.eval()
    onnx_model = onnx.load( outModel )
    providers = ['CPUExecutionProvider']
    so = onnxruntime.SessionOptions()
    so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), so, providers=providers)

    wavs = [
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/multi-speaker_4-speakers_Jennifer_Aniston_and_Adam_Sandler_talk.wav',
            #'/home/leo/storage/sharedFolderVirtualbox/audioForTesting/multi-speaker_1min.wav',
            #'/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching1.wav',
            #'/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav',
            #'/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav',
            ]


    batch = 32
    wav_lens1 = [1.0] * batch
    wav_lens1 = torch.Tensor( wav_lens1 )
    wav_lens1 = wav_lens1.to(torch.float32)
    for wav in wavs:
        print( "\n-----------------------------" )
        # Load wav
        signal, fs = torchaudio.load( wav )
        signal = signal[:,:batch*80000].flatten().reshape( batch, -1 )

        # Calc embeddings using pytorch
        print( signal.shape )
        compute_STFT = STFT(sample_rate=16000, win_length=25, hop_length=10, n_fft=400)
        feats_stft = compute_STFT(signal)
        np.savetxt(f's4_feats.txt', feats_stft.flatten(), fmt='%f', delimiter=',')

        # Calc embeddings using pytorch
        print( feats_stft.shape )
        #wav_lens = torch.Tensor([1.0])
        embeddings_origin = model_origin.encode_batch(signal, normalize=False)
        np.savetxt(f's4_emd.txt', embeddings_origin.flatten(), fmt='%f', delimiter=',')
        print( embeddings_origin  )
        embeddings = model.encode_batch( feats_stft, wav_lens1 )

        # Calc embedding using onnx model
        print( ort_session.get_inputs()[0])
        print( ort_session.get_inputs()[1])
        ort_inputs = {ort_session.get_inputs()[0].name: feats_stft.numpy(),
                ort_session.get_inputs()[1].name: wav_lens1.numpy()
                }
        ort_outs = ort_session.run(None, ort_inputs)
        onnx_embeddings = ort_outs[0]
        print(onnx_embeddings)

        try:
            # Check if result is close enough
            np.testing.assert_allclose(
                    embeddings.detach().numpy(), 
                    onnx_embeddings, 
                    rtol=1e-03, 
                    atol=1e-05,
                    verbose = True)
            print( f'check passed for {wav}' )
        except AssertionError as e:
            print( e )
        print( "-----------------------------\n" )

def test():
    from speechbrain.processing.features import STFT
    model = STFT(sample_rate=16000, win_length=25, hop_length=10, n_fft=400)
    #model = MySTFT(sample_rate=16000, win_length=25, hop_length=10, n_fft=400)

    wavs = [
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/multi-speaker_1min.wav',
            #'/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching1.wav',
            #'/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav',
            ]

    i = 0
    for wav in wavs:
        # Load wav
        signal, fs = torchaudio.load( wav )
        #signal = signal[:,:80000]
        signal = signal[:,:400]
        #signal = signal[:,:800]
        #signal = signal[:,:320000]
        #signal = signal.flatten().reshape( 4, -1 )
        print( signal.shape )
        feats = model(signal)
        print('--------------------')
        print( feats.shape )
        #print(feats)
        np.savetxt(f'/tmp/stft_wav{i}.txt', signal.detach().numpy().flatten(), fmt='%f')
        np.savetxt(f'/tmp/stft_feats{i}.txt', feats.detach().numpy().flatten(), fmt='%f')
        i += 1

def manualFbank():
    import kaldi_native_fbank as knf
    model = FeatureModel.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")
    model.eval()

    wavs = [
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching1.wav',
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav',
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav',
            ]

    for wav in wavs:
        print( f"Check wav file: {wav}" )

        # Load wav
        signal, fs = torchaudio.load( wav )
        signal = signal[:,:80000]

        # Calc embeddings using pytorch
        print( signal.shape )
        wav_lens = torch.Tensor([1.0])
        feats = model.encode_batch(signal)
        feats.detach().numpy(), 

        # Calc using fbank 
        sampling_rate = 16000
        opts = knf.FbankOptions()
        opts.frame_opts.dither = 0
        opts.frame_opts.remove_dc_offset = False
        opts.frame_opts.snip_edges = False
        opts.mel_opts.num_bins = 80
        opts.mel_opts.debug_mel = False
        opts.mel_opts.low_freq = 0
        opts.mel_opts.high_freq = 8000
        fbank = knf.OnlineFbank(opts)
        fbank.accept_waveform(sampling_rate, signal[0].tolist())
        ffeats = []
        for i in range(fbank.num_frames_ready):
            f = torch.from_numpy(fbank.get_frame(i))
            ffeats.append( f )


        try:
            # Check if result is close enough
            np.testing.assert_allclose(
                    feats[0][0], 
                    ffeats[0], 
                    rtol=1e-03, 
                    atol=1e-05,
                    verbose = True)
            print( 'check passed' )
        except AssertionError as e:
            print( e )
        print( "-----------------------------" )

def verify_stft( outModel ):
    from speechbrain.processing.features import STFT
    model = STFT(sample_rate=16000, win_length=25, hop_length=10, n_fft=400)

    onnx_model = onnx.load( outModel )
    providers = ['CPUExecutionProvider']
    so = onnxruntime.SessionOptions()
    so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), so, providers=providers)

    wavs = [
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching1.wav',
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav',
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav',
            ]

    for wav in wavs:
        print( f"Check wav file: {wav}" )

        # Load wav
        signal, fs = torchaudio.load( wav )
        signal = signal[:,:80000]

        # Calc embeddings using pytorch
        print( signal.shape )
        feats = model(signal)

        # Calc embedding using onnx model
        print( ort_session.get_inputs()[0])
        ort_inputs = {ort_session.get_inputs()[0].name: signal.numpy(),
                }
        ort_outs = ort_session.run(None, ort_inputs)
        onnx_feats = ort_outs[0]

        try:
            # Check if result is close enough
            np.testing.assert_allclose(
                    feats.detach().numpy(), 
                    onnx_feats, 
                    rtol=1e-03, 
                    atol=1e-05,
                    verbose = True)
            print( 'check passed' )
        except AssertionError as e:
            print( e )
        print( "-----------------------------" )

outModel = 'feature3.onnx'
#export(outModel )
#debug(outModel )
#verify(outModel )

'''
input: features 
output: embedding
This model works
'''
outModel = 'emd3.onnx'
#export1(outModel )
#verify1(outModel )

'''
input: output of stft 
output: embedding
This model works
'''
outModel = 'emd4.onnx'
export2(outModel )
verify2(outModel )

'''
This model does not work
'''
outModel = 'stft.onnx'
#export_stft( outModel )
#verify_stft( outModel )

'''
This model works
'''
outModel = 'fbank.onnx'
#export_fbank( outModel )
#verify_fbank( outModel )

# Calc fbank features using kaldi-native-fbank
# https://github.com/csukuangfj/kaldi-native-fbank
#manualFbank()

#test()
#simple()

#genFeats()

embeddings/export3.py
leohuang2013/pyannote-audio_speaker-diarization_cpp

# put your own hugging face token
hf_auth_token="YOUR_OWN_HF_TOKEN"

embeddings/hf_token.py
leohuang2013/pyannote-audio_speaker-diarization_cpp

import torchaudio
from speechbrain.pretrained import EncoderClassifier
import torch
from hyperpyyaml import load_hyperpyyaml


ckpt = "pretrained_models/EncoderClassifier-8f6f7fdaa9628acf73e21ad1f99d5f83/hyperparams.yaml"

hparams = load_hyperpyyaml(
    open( ckpt ),
)

model = EncoderClassifier(modules=hparams['modules'], hparams=dict(hparams))

model.eval()

x = torch.randn( 1, 1920000 )
print( x )
print( x.shape )
#exit( 0 )

traced_script_module = torch.jit.trace(model, x)
traced_script_module.save("traced_model.pt")


embeddings/jitload.py
leohuang2013/pyannote-audio_speaker-diarization_cpp
import torch
import torchaudio
from myembedding import MyEmbedding
from speechbrain.pretrained import EncoderClassifier
import numpy as np

outModel = 'embedding.pt'
def export():
    # Create dummy input
    # Batch of waveforms [batch, time, channels] or [batch, time]
    # speechbrain/pretrained/interfaces.py:903
    #batch = 32
    batch = 1
    max_sequence = 80000
    #x = torch.FloatTensor(batch, max_sequence).uniform_(-32768, 32767)
    x = torch.FloatTensor(batch, max_sequence).uniform_(-1, 1)
    print( x.shape )
    #exit( 0 ) 

    model = MyEmbedding.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")
    model.eval()
    pretrained_dict = model.state_dict()
    #print( pretrained_dict )
    #exit(0)

    signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav')
    print( '\n---- start export ----' )
    # Export the model
    '''
    There will be big difference between using signal and x for result of mean_var_norm
    '''
    #model_script = torch.jit.script( model )
    wav_lens = torch.Tensor([1.0])
    wav_lens1 = [1.0] * batch
    wav_lens1 = torch.Tensor( wav_lens1 )
    wav_lens1 = wav_lens1.to(torch.float32)
    traced_model = torch.jit.trace(model, [x, wav_lens1])
    torch.jit.save(traced_model, outModel)
    print( f'---- model has been saved to {outModel} ----\n' )

def verify():
    model = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="pretrained")
    loaded_model = torch.jit.load( outModel )
    loaded_model.eval()
    loaded_model.to(torch.device('cuda'))

    wavs = [
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching1.wav',
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav',
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav',
            ]

    for wav in wavs:
        print( f"Check wav file: {wav}" )

        # Load wav
        signal, fs = torchaudio.load( wav )
        signal = signal[:80000]

        # Calc embeddings using pytorch
        print( signal.shape )
        wav_lens = torch.Tensor([1.0])
        embeddings = model.encode_batch(signal, wav_lens, normalize=False)

        # Calc embedding using onnx model
        pt_embeddings = loaded_model(signal, wav_lens)

        try:
            # Check if result is close enough
            np.testing.assert_allclose(
                    embeddings.detach().numpy(), 
                    pt_embeddings, 
                    rtol=1e-03, 
                    atol=1e-05,
                    verbose = True)
            print( 'check passed' )
        except AssertionError as e:
            print( e )
        print( "-----------------------------" )

#export()
verify()

embeddings/mm.py
leohuang2013/pyannote-audio_speaker-diarization_cpp

import torch
from speechbrain.pretrained import Pretrained

class MyEmbedding(Pretrained):
    """A ready-to-use class for utterance-level classification (e.g, speaker-id,
    language-id, emotion recognition, keyword spotting, etc).

    The class assumes that an encoder called "embedding_model" and a model
    called "classifier" are defined in the yaml file. If you want to
    convert the predicted index into a corresponding text label, please
    provide the path of the label_encoder in a variable called 'lab_encoder_file'
    within the yaml.

    The class can be used either to run only the encoder (encode_batch()) to
    extract embeddings or to run a classification step (classify_batch()).
    ```

    Example
    -------
    >>> import torchaudio
    >>> from speechbrain.pretrained import EncoderClassifier
    >>> # Model is downloaded from the speechbrain HuggingFace repo
    >>> tmpdir = getfixture("tmpdir")
    >>> classifier = EncoderClassifier.from_hparams(
    ...     source="speechbrain/spkrec-ecapa-voxceleb",
    ...     savedir=tmpdir,
    ... )

    >>> # Compute embeddings
    >>> signal, fs = torchaudio.load("tests/samples/single-mic/example1.wav")
    >>> embeddings = classifier.encode_batch(signal)

    >>> # Classification
    >>> prediction = classifier.classify_batch(signal)
    """

    MODULES_NEEDED = [
        "compute_features",
        "mean_var_norm",
        "embedding_model",
        "classifier",
    ]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def encode_batch(self, wavs, wav_lens=None, normalize=False):
        """Encodes the input audio into a single vector embedding.

        The waveforms should already be in the model's desired format.
        You can call:
        ``normalized = <this>.normalizer(signal, sample_rate)``
        to get a correctly converted signal in most cases.

        Arguments
        ---------
        wavs : torch.Tensor
            Batch of waveforms [batch, time, channels] or [batch, time]
            depending on the model. Make sure the sample rate is fs=16000 Hz.
        wav_lens : torch.Tensor
            Lengths of the waveforms relative to the longest one in the
            batch, tensor of shape [batch]. The longest one should have
            relative length 1.0 and others len(waveform) / max_length.
            Used for ignoring padding.
        normalize : bool
            If True, it normalizes the embeddings with the statistics
            contained in mean_var_norm_emb.

        Returns
        -------
        torch.Tensor
            The encoded batch
        """
        # Manage single waveforms in input
        #if len(wavs.shape) == 1:
        if len(wavs.size()) == 1:
            wavs = wavs.unsqueeze(0)

        # Assign full length if wav_lens is not assigned
        if wav_lens is None:
            #wav_lens = torch.ones(wavs.shape[0], device=self.device)
            wav_lens = torch.ones(wavs.size(0), device=self.device)

        # Storing waveform in the specified device
        wavs, wav_lens = wavs.to(self.device), wav_lens.to(self.device)
        wavs = wavs.float()

        # Computing features and embeddings
        self.mods.to( self.device )
        feats = self.mods.compute_features(wavs)
        #feats = self.mods.mean_var_norm(feats, wav_lens)
        mn = MyNormalization()
        feats = mn.forward( feats, wav_lens )
        embeddings = self.mods.embedding_model(feats, wav_lens)
        if normalize:
            embeddings = self.hparams.mean_var_norm_emb(
                #embeddings, torch.ones(embeddings.shape[0], device=self.device)
                embeddings, torch.ones(embeddings.size(0), device=self.device)
            )
        return embeddings

    def forward(self, wavs, wav_lens=None):
        """Runs the classification"""
        # LIYI
        #return self.classify_batch(wavs, wav_lens)
        return self.encode_batch(wavs, wav_lens)
        #feats = wavs
        #return self.mods.embedding_model(feats, wav_lens)


class MyNormalization:
    """Performs mean and variance normalization of the input tensor.

    Arguments
    ---------
    mean_norm : True
         If True, the mean will be normalized.
    std_norm : True
         If True, the standard deviation will be normalized.
    norm_type : str
         It defines how the statistics are computed ('sentence' computes them
         at sentence level, 'batch' at batch level, 'speaker' at speaker
         level, while global computes a single normalization vector for all
         the sentences in the dataset). Speaker and global statistics are
         computed with a moving average approach.
    avg_factor : float
         It can be used to manually set the weighting factor between
         current statistics and accumulated ones.

    Example
    -------
    >>> import torch
    >>> norm = InputNormalization()
    >>> inputs = torch.randn([10, 101, 20])
    >>> inp_len = torch.ones([10])
    >>> features = norm(inputs, inp_len)
    """


    def __init__(
        self,
        mean_norm=True,
        std_norm=False,
        norm_type="sentence",
    ):
        super().__init__()
        self.mean_norm = mean_norm
        self.std_norm = std_norm
        self.norm_type = norm_type
        self.eps = 1e-10

    def forward(self, x, lengths):
        """Returns the tensor with the surrounding context.

        Arguments
        ---------
        x : tensor
            A batch of tensors.
        lengths : tensor
            A batch of tensors containing the relative length of each
            sentence (e.g, [0.7, 0.9, 1.0]). It is used to avoid
            computing stats on zero-padded steps.
        spk_ids : tensor containing the ids of each speaker (e.g, [0 10 6]).
            It is used to perform per-speaker normalization when
            norm_type='speaker'.
        """
        #N_batches = x.shape[0]
        N_batches = x.size(0)
        seq_len = x.size(1)

        for snt_id in range(N_batches):

            # Avoiding padded time steps
            # Changed to int64, otherwise go error:
            # Type parameter (Tind) of Optype (Slice) bound to different types (tensor(int64) and tensor(int32) in node (/Slice)
            #actual_size = torch.round(lengths[snt_id] * x.shape[1]).long()
            actual_size = torch.round(lengths[snt_id] * seq_len).long()

            # computing statistics
            current_mean, current_std = self._compute_current_stats(
                x[snt_id, 0:actual_size, ...]
            )

            if self.norm_type == "sentence":
                x[snt_id] = (x[snt_id] - current_mean) / current_std


        return x

    def _compute_current_stats(self, x):
        """Returns the tensor with the surrounding context.

        Arguments
        ---------
        x : tensor
            A batch of tensors.
        """
        # Compute current mean
        if self.mean_norm:
            current_mean = torch.mean(x, dim=0).detach()
        else:
            current_mean = torch.tensor([0.0], device=x.device)

        # Compute current std
        if self.std_norm:
            current_std = torch.std(x, dim=0).detach()
        else:
            current_std = torch.tensor([1.0], device=x.device)

        # Improving numerical stability of std
        current_std = torch.max(
            current_std, self.eps * torch.ones_like(current_std)
        )

        return current_mean, current_std

embeddings/myembedding.py
leohuang2013/pyannote-audio_speaker-diarization_cpp
# 1. visit hf.co/pyannote/embedding and accept user conditions
# 2. visit hf.co/settings/tokens to create an access token
# 3. instantiate pretrained model
from pyannote.audio import Model
import torch

hf_auth_token="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"
model = Model.from_pretrained("pyannote/embedding", 
                              use_auth_token=hf_auth_token)

from pyannote.audio import Inference

signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav')
print( signal )
print( signal.shape )
#exit( 0 ) 

inference = Inference(model, window="whole")
embedding1 = inference("/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav")
# `embeddingX` is (1 x D) numpy array extracted from the file as a whole.

# `distance` is a `float` describing how dissimilar speakers 1 and 2 are.

symbolic_names = {0: "batch_size", 1: "max_seq_len"}
x = torch.randn( 1, 1920000 )
print( x )
print( x.shape )
#exit( 0 ) 
torch.onnx.export(model,               # model being run
                  x,                         # model input (or a tuple for multiple inputs)
                  "embeddings.onnx",   # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=17,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  verbose=False,
                  input_names = ['signal'],   # the model's input names
                  output_names = ['embeddings'], # the model's output names
                  dynamic_axes={'signal' : symbolic_names,    # variable length axes
                                })

embeddings/pyannote-embeddings.py
leohuang2013/pyannote-audio_speaker-diarization_cpp
import torch
import torch.onnx.verification
torch.manual_seed(0)
opset_version = 15

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(3, 4),
            torch.nn.ReLU(),
            torch.nn.Linear(4, 5),
            torch.nn.ReLU(),
            torch.nn.Linear(5, 6),
        )
    def forward(self, x):
        return self.layers(x)

model = Model()
graph_info = torch.onnx.verification.find_mismatch(
    model,
    (torch.randn(2, 3),),
    opset_version=opset_version,
)

torch.onnx.export(
    model,
    (torch.randn(2, 3),),
    'simple.onnx',
    opset_version=opset_version,
    do_constant_folding = True,
    verbose = True
)

print("\n=============")
print("exported")

embeddings/simple.py
leohuang2013/pyannote-audio_speaker-diarization_cpp

import torch
from speechbrain.pretrained import Pretrained

from speechbrain.processing.features import Filterbank, spectral_magnitude

class MySTFT(torch.nn.Module):
    def __init__(
        self,
        sample_rate,
        win_length=25,
        hop_length=10,
        n_fft=400,
        window_fn=torch.hamming_window,
        normalized_stft=False,
        center=True,
        pad_mode="constant",
        onesided=True,
    ):
        super().__init__()
        self.sample_rate = sample_rate
        self.win_length = win_length
        self.hop_length = hop_length
        self.n_fft = n_fft
        self.normalized_stft = normalized_stft
        self.center = center
        self.pad_mode = pad_mode
        self.onesided = onesided

        # Convert win_length and hop_length from ms to samples
        self.win_length = int(
            round((self.sample_rate / 1000.0) * self.win_length)
        )
        self.hop_length = int(
            round((self.sample_rate / 1000.0) * self.hop_length)
        )

        self.window = window_fn(self.win_length)


    def forward(self, x):
        """Returns the STFT generated from the input waveforms.

        Arguments
        ---------
        x : tensor
            A batch of audio signals to transform.
        """

        stft = torch.stft(
            x,
            self.n_fft,
            self.hop_length,
            self.win_length,
            self.window.to(x.device),
            self.center,
            self.pad_mode,
            self.normalized_stft,
            self.onesided,
            return_complex=False,
        )

        # Retrieving the original dimensionality (batch,time, channels)
        stft = stft.transpose(2, 1)

        return stft


class FBankModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        features = spectral_magnitude(x)
        compute_fbanks = Filterbank(n_mels=80)
        feats = compute_fbanks(features)
        return feats

class FeatureModel(Pretrained):
    MODULES_NEEDED = [
        "compute_features",
        "mean_var_norm",
        "embedding_model",
        "classifier",
    ]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def encode_batch(self, wavs, wav_lens=None, normalize=False ):
        """Encodes the input audio into a single vector embedding.

        The waveforms should already be in the model's desired format.
        You can call:
        ``normalized = <this>.normalizer(signal, sample_rate)``
        to get a correctly converted signal in most cases.

        Arguments
        ---------
        wavs : torch.Tensor
            Batch of waveforms [batch, time, channels] or [batch, time]
            depending on the model. Make sure the sample rate is fs=16000 Hz.
        wav_lens : torch.Tensor
            Lengths of the waveforms relative to the longest one in the
            batch, tensor of shape [batch]. The longest one should have
            relative length 1.0 and others len(waveform) / max_length.
            Used for ignoring padding.
        normalize : bool
            If True, it normalizes the embeddings with the statistics
            contained in mean_var_norm_emb.

        Returns
        -------
        torch.Tensor
            The encoded batch
        """
        # Storing waveform in the specified device
        wavs = wavs.to(self.device)
        wavs = wavs.float()

        # Computing features and embeddings
        self.mods.to( self.device )
        from speechbrain.processing.features import Filterbank, STFT
        from speechbrain.processing.features import spectral_magnitude
        compute_STFT = STFT(sample_rate=16000, win_length=25, hop_length=10, n_fft=400)
        features = compute_STFT(wavs)
        features = spectral_magnitude(features)
        compute_fbanks = Filterbank(n_mels=80)
        feats = compute_fbanks(features)
        #feats = self.mods.compute_features(wavs)
        return feats

    def forward(self, wavs, wav_lens=None):
        """Runs the classification"""
        # LIYI
        #return self.classify_batch(wavs, wav_lens)
        return self.encode_batch(wavs, wav_lens)
        #feats = wavs
        #return self.mods.embedding_model(feats, wav_lens)

class MyEmbedding0(Pretrained):
    """A ready-to-use class for utterance-level classification (e.g, speaker-id,
    language-id, emotion recognition, keyword spotting, etc).

    The class assumes that an encoder called "embedding_model" and a model
    called "classifier" are defined in the yaml file. If you want to
    convert the predicted index into a corresponding text label, please
    provide the path of the label_encoder in a variable called 'lab_encoder_file'
    within the yaml.

    The class can be used either to run only the encoder (encode_batch()) to
    extract embeddings or to run a classification step (classify_batch()).
    ```

    Example
    -------
    >>> import torchaudio
    >>> from speechbrain.pretrained import EncoderClassifier
    >>> # Model is downloaded from the speechbrain HuggingFace repo
    >>> tmpdir = getfixture("tmpdir")
    >>> classifier = EncoderClassifier.from_hparams(
    ...     source="speechbrain/spkrec-ecapa-voxceleb",
    ...     savedir=tmpdir,
    ... )

    >>> # Compute embeddings
    >>> signal, fs = torchaudio.load("tests/samples/single-mic/example1.wav")
    >>> embeddings = classifier.encode_batch(signal)

    >>> # Classification
    >>> prediction = classifier.classify_batch(signal)
    """

    MODULES_NEEDED = [
        "compute_features",
        "embedding_model",
    ]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def encode_batch(self, stft_out, wav_lens=None, normalize=False):
        """Encodes the input audio into a single vector embedding.

        The waveforms should already be in the model's desired format.
        You can call:
        ``normalized = <this>.normalizer(signal, sample_rate)``
        to get a correctly converted signal in most cases.

        Arguments
        ---------
        wavs : torch.Tensor
            Batch of waveforms [batch, time, channels] or [batch, time]
            depending on the model. Make sure the sample rate is fs=16000 Hz.
        wav_lens : torch.Tensor
            Lengths of the waveforms relative to the longest one in the
            batch, tensor of shape [batch]. The longest one should have
            relative length 1.0 and others len(waveform) / max_length.
            Used for ignoring padding.
        normalize : bool
            If True, it normalizes the embeddings with the statistics
            contained in mean_var_norm_emb.

        Returns
        -------
        torch.Tensor
            The encoded batch
        """

        # Storing waveform in the specified device
        wav_lens = wav_lens.to(self.device)

        features = spectral_magnitude(stft_out)
        compute_fbanks = Filterbank(n_mels=80)
        feats = compute_fbanks(features)

        # Computing features and embeddings
        self.mods.to( self.device )
        #feats = self.mods.mean_var_norm(feats, wav_lens)
        mn = MyNormalization()
        feats = mn.forward( feats, wav_lens )
        embeddings = self.mods.embedding_model(feats, wav_lens)
        if normalize:
            embeddings = self.hparams.mean_var_norm_emb(
                #embeddings, torch.ones(embeddings.shape[0], device=self.device)
                embeddings, torch.ones(embeddings.size(0), device=self.device)
            )
        return embeddings

    def forward(self, stft_out, wav_lens):
        """Runs the classification"""
        # LIYI
        return self.encode_batch(stft_out, wav_lens)


class EMDModel(Pretrained):
    MODULES_NEEDED = [
        "compute_features",
        "mean_var_norm",
        "embedding_model",
        "classifier",
    ]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def encode_batch(self, feats, wav_lens=None):
        """Encodes the input audio into a single vector embedding.

        The waveforms should already be in the model's desired format.
        You can call:
        ``normalized = <this>.normalizer(signal, sample_rate)``
        to get a correctly converted signal in most cases.

        Arguments
        ---------
        wavs : torch.Tensor
            Batch of waveforms [batch, time, channels] or [batch, time]
            depending on the model. Make sure the sample rate is fs=16000 Hz.
        wav_lens : torch.Tensor
            Lengths of the waveforms relative to the longest one in the
            batch, tensor of shape [batch]. The longest one should have
            relative length 1.0 and others len(waveform) / max_length.
            Used for ignoring padding.
        normalize : bool
            If True, it normalizes the embeddings with the statistics
            contained in mean_var_norm_emb.

        Returns
        -------
        torch.Tensor
            The encoded batch
        """
        # Assign full length if wav_lens is not assigned
        if wav_lens is None:
            #wav_lens = torch.ones(wavs.shape[0], device=self.device)
            wav_lens = torch.ones(feats.size(0), device=self.device)

        # Storing waveform in the specified device
        feats, wav_lens = feats.to(self.device), wav_lens.to(self.device)
        feats = feats.float()

        mn = MyNormalization()
        feats = mn.forward( feats, wav_lens )
        self.mods.to( self.device )
        embeddings = self.mods.embedding_model(feats, wav_lens)
        return embeddings

    def forward(self, feats, wav_lens):
        """Runs the classification"""
        return self.encode_batch(feats, wav_lens)

class MyNormalization:
    """Performs mean and variance normalization of the input tensor.

    Arguments
    ---------
    mean_norm : True
         If True, the mean will be normalized.
    std_norm : True
         If True, the standard deviation will be normalized.
    norm_type : str
         It defines how the statistics are computed ('sentence' computes them
         at sentence level, 'batch' at batch level, 'speaker' at speaker
         level, while global computes a single normalization vector for all
         the sentences in the dataset). Speaker and global statistics are
         computed with a moving average approach.
    avg_factor : float
         It can be used to manually set the weighting factor between
         current statistics and accumulated ones.

    Example
    -------
    >>> import torch
    >>> norm = InputNormalization()
    >>> inputs = torch.randn([10, 101, 20])
    >>> inp_len = torch.ones([10])
    >>> features = norm(inputs, inp_len)
    """


    def __init__(
        self,
        mean_norm=True,
        std_norm=False,
        norm_type="sentence",
    ):
        super().__init__()
        self.mean_norm = mean_norm
        self.std_norm = std_norm
        self.norm_type = norm_type
        self.eps = 1e-10

    def forward(self, x, lengths):
        """Returns the tensor with the surrounding context.

        Arguments
        ---------
        x : tensor
            A batch of tensors.
        lengths : tensor
            A batch of tensors containing the relative length of each
            sentence (e.g, [0.7, 0.9, 1.0]). It is used to avoid
            computing stats on zero-padded steps.
        spk_ids : tensor containing the ids of each speaker (e.g, [0 10 6]).
            It is used to perform per-speaker normalization when
            norm_type='speaker'.
        """
        #N_batches = x.shape[0]
        N_batches = x.size(0)
        seq_len = x.size(1)

        for snt_id in range(N_batches):

            # Avoiding padded time steps
            # Changed to int64, otherwise go error:
            # Type parameter (Tind) of Optype (Slice) bound to different types (tensor(int64) and tensor(int32) in node (/Slice)
            #actual_size = torch.round(lengths[snt_id] * x.shape[1]).long()
            actual_size = torch.round(lengths[snt_id] * seq_len).long()

            # computing statistics
            current_mean, current_std = self._compute_current_stats(
                x[snt_id, 0:actual_size, ...]
            )

            if self.norm_type == "sentence":
                x[snt_id] = (x[snt_id] - current_mean) / current_std


        return x

    def _compute_current_stats(self, x):
        """Returns the tensor with the surrounding context.

        Arguments
        ---------
        x : tensor
            A batch of tensors.
        """
        # Compute current mean
        if self.mean_norm:
            current_mean = torch.mean(x, dim=0).detach()
        else:
            current_mean = torch.tensor([0.0], device=x.device)

        # Compute current std
        if self.std_norm:
            current_std = torch.std(x, dim=0).detach()
        else:
            current_std = torch.tensor([1.0], device=x.device)

        # Improving numerical stability of std
        current_std = torch.max(
            current_std, self.eps * torch.ones_like(current_std)
        )

        return current_mean, current_std

embeddings/threeModel.py
leohuang2013/pyannote-audio_speaker-diarization_cpp
import os
import sys
import numpy as np

#checkList = ['batch_waveform','segmentations','batch_masks', 'wav_lens', 'signals', 'masks',
checkList = ['batch_masks', 'segmentations', 'wav_lens', 'signals', 'masks',
        'imasks', 'binarize_score', 'on', 'same_as', 'samples', 'well_defined_idx', 
        'initial_state', 'binarized_segmentations', 
        'binary_ndarray', 'final_wav_lens', 
        'clean_segmentations', 'trimmed', 'sum_trimmed', 'count_data',
        'count', 'embeddings','filtered_embeddings',
        'norm_embeddings', 'clusters', 'dist', 'clusterRes', 
        'soft_clusters', 'hard_clusters', 'clustered_segmentations',
        'aggregated_output', 'aggregated_mask', 'overlapping_chunk_count',
        'scores_in_aggregate', 'masks_in_aggregate',
        'to_diarization_activations', 'cropped_activations','cropped_count', 
        'sorted_speakers', 'discrete_diarization' ]

def deleteSingle( item, source ):
        more = True
        num = 0
        while more:
            fn = f'/tmp/{source}_{item}{num}.txt'
            if os.path.exists( fn ):
                print( f'Delete file: {fn}' )
                os.remove( fn )
                num += 1
            else:
                more = False

        fn = f'/tmp/{source}_{item}.txt'
        if os.path.exists( fn ):
            print( f'Delete file: {fn}' )
            os.remove( fn )

def deleteResultFiles():
    for item in checkList:
        deleteSingle( item, 'cpp' )
        deleteSingle( item, 'py' )

def checkIfSame( fn_cpp, fn_py ):
    f_cpp = open( fn_cpp )
    f_py = open( fn_py )
    if f_cpp.read() == f_py.read():
        same = True
    else:
        same = False
    f_cpp.close()
    f_py.close()
    return same

def requireSameFileContent( item ):
    more = True
    num = 0
    while more:
        fn_cpp = f'/tmp/cpp_{item}{num}.txt'
        fn_py = f'/tmp/py_{item}{num}.txt'
        if os.path.exists( fn_cpp ) and os.path.exists( fn_py ):
            if checkIfSame( fn_cpp, fn_py ):
                print( f"Checking passed for {item}{num}" )
            else:
                print( f"Difference is detected for {item}{num}" )
            num += 1
        else:
            more = False

    fn_cpp = f'/tmp/cpp_{item}.txt'
    fn_py = f'/tmp/py_{item}.txt'
    if os.path.exists( fn_cpp ) and os.path.exists( fn_py ):
        if checkIfSame( fn_cpp, fn_py ):
            print( f"Checking passed for {item}" )
        else:
            print( f"Difference is detected for {item}" )

def local_check( cpp_arr, py_arr ):
    import math
    for i in range( len( cpp_arr )):
        if cpp_arr[i] == py_arr[i]:
            continue
        elif math.isnan( cpp_arr[i] ) and math.isnan( py_arr[i] ):
            continue
        else:
            if math.isnan( cpp_arr[i] ) or math.isnan( py_arr[i] ):
                print( 'nan is different' )
            else:
                print( f'[{i}]: float different: ' + str( cpp_arr[i] - py_arr[i]) )


def checkCloseEnough( fn_cpp, fn_py ):
    f_cpp = open( fn_cpp )
    f_py = open( fn_py )
    cpp_content = f_cpp.read()
    py_content = f_py.read()
    py_content = py_content.replace( '\n', ',' )
    cpp_content = cpp_content.replace( '\n', ',' )
    pyarr = py_content.split( ',' )
    cpparr = cpp_content.split( ',' )
    py_arr = []
    cpp_arr = []
    for a in pyarr:
        if a == '' or a == ' ':
            continue
        py_arr.append( float( a ))
    for a in cpparr:
        if a == '' or a == ' ':
            continue
        cpp_arr.append( float( a ))
    same = True
    #if 'soft_clusters' in fn_cpp:
    #    local_check( cpp_arr, py_arr )
    if len( py_arr ) != len( cpp_arr ):
        print( f'size is different, py: {len(py_arr)}, cpp: {len(cpp_arr)}' )
        return False
    else:
        py_nparr = np.array( py_arr )
        cpp_nparr = np.array( cpp_arr )
        try:
            # Check if result is close enough
            np.testing.assert_allclose(
                    py_nparr,
                    cpp_nparr,
                    rtol=0.001,
                    atol=0.0001,
                    verbose = True)
            print( 'check passed' )
        except AssertionError as e:
            same = False
            print( e )
    f_cpp.close()
    f_py.close()
    if not same:
        print( '\n----------------- summary ------------------' )
    return same

def requireCloseEnough( item ):
    more = True
    num = 0
    while more:
        fn_cpp = f'/tmp/cpp_{item}{num}.txt'
        fn_py = f'/tmp/py_{item}{num}.txt'
        if os.path.exists( fn_cpp ) and os.path.exists( fn_py ):
            if checkCloseEnough( fn_cpp, fn_py ):
                print( f"Checking passed for {item}{num}" )
            else:
                print( f"Difference is detected for {item}{num}" )
            num += 1
        else:
            more = False

    fn_cpp = f'/tmp/cpp_{item}.txt'
    fn_py = f'/tmp/py_{item}.txt'
    if os.path.exists( fn_cpp ) and os.path.exists( fn_py ):
        if checkCloseEnough( fn_cpp, fn_py ):
            print( f"Checking passed for {item}" )
        else:
            print( f"Difference is detected for {item}" )

def checkEmbeddings():
    pass

def checkResult():
    sameFileContentList = ['same_as', 'well_defined_idx', 'samples', 'on', 'initial_state', 
            'masks', 'imasks', 'wav_lens', 'signals', 'count', 'clusters']
    closeEnoughList = ['segmentations', 'clean_segmentations', 'binarize_score', 
            'binarized_segmentations', 
            'trimmed', 'sum_trimmed', 'count_data', 'batch_waveform', 'embeddings',
            'norm_embeddings', 'dist', 'clusterRes', 'soft_clusters', 'hard_clusters',
            'clustered_segmentations', 'filtered_embeddings','aggregated_output',
            'aggregated_mask', 'overlapping_chunk_count', 'masks_in_aggregate',
            'scores_in_aggregate', 'to_diarization_activations',
            'cropped_activations', 'cropped_count','sorted_speakers', 'discrete_diarization']
    for item in checkList:
        print(f'\n======================================')
        print(f'======== CHECKING {item} =========')
        if item in sameFileContentList:
            requireSameFileContent( item )
        elif item in closeEnoughList:
            requireCloseEnough( item )

def main():
    if len( sys.argv ) > 1 and sys.argv[1] == 'clean':
        deleteResultFiles()
        return
    else:
        checkResult()


if __name__ == '__main__':
    main()

pipeline/script/verifyEveryStepResult.py
leohuang2013/pyannote-audio_speaker-diarization_cpp

import torchaudio
import torch
import numpy as np
from mysegment import MySegment

opsetVer = 17
outModel = 'segment.onnx'

def export():
    # Create dummy input
    #audio = '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/shortTeaching2.wav'
    audio = '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/multi-speaker_4-speakers_Jennifer_Aniston_and_Adam_Sandler_talk.wav'
    signal, fs = torchaudio.load( audio )
    print( signal.shape )

    self.model =  Model.from_pretrained(
            "pyannote/segmentation@2022.07",
            strict=False,
            use_auth_token=hf_auth_token,
        )
    self.model.eval()
    self.model.to( torch.device( 'cuda' ))
    model = MySegment()

    # Show all unconvertable ops, output would be like,
    #       {'aten::view_as_real', 'aten::chunk', 'aten::stft'}
    #torch_script_graph, unconvertible_ops = torch.onnx.utils.unconvertible_ops(
    #    model, signal, opset_version=opsetVer
    #)
    #if( len( unconvertible_ops ) > 0 ):
    #    print( '------------------------There are some unconvertable operations-----------' )
    #    print(set(unconvertible_ops))
    #    exit( 0 )
    #print( '---- all operations convertable ----' )

    # Export the model
    print( '\n---- start export ----' )
    symbolic_names = {0: "batch_size", 1: "max_seq_len"}
    torch.onnx.export(model,               # model being run
                      signal,                         # model input (or a tuple for multiple inputs)
                      outModel,   # where to save the model (can be a file or file-like object)
                      export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=opsetVer,          # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      verbose=False,
                      input_names = ['signal'],   # the model's input names
                      output_names = ['segments'], # the model's output names
                      dynamic_axes={'signal' : symbolic_names,    # variable length axes
                                    })
    print( f'---- model has been saved to {outModel} ----\n' )


export()

segment/export.py
leohuang2013/pyannote-audio_speaker-diarization_cpp

import torchaudio
import torch
import numpy as np
from pyannote.audio.core.model import Model
from hf_token import *
import torchaudio
import onnx
import onnxruntime

opsetVer = 17
outModel = 'segment2.onnx'
batch = 32
channel = 1

def export():
    model =  Model.from_pretrained(
            "pyannote/segmentation@2022.07",
            strict=False,
            use_auth_token=hf_auth_token,
        )
    model.eval()

    # Create dummy input
    dummy_input = torch.zeros(batch, channel, 80000)

    # Show all unconvertable ops, output would be like,
    #       {'aten::view_as_real', 'aten::chunk', 'aten::stft'}
    #torch_script_graph, unconvertible_ops = torch.onnx.utils.unconvertible_ops(
    #    model, signal, opset_version=opsetVer
    #)
    #if( len( unconvertible_ops ) > 0 ):
    #    print( '------------------------There are some unconvertable operations-----------' )
    #    print(set(unconvertible_ops))
    #    exit( 0 )
    #print( '---- all operations convertable ----' )

    # Export the model
    print( '\n---- start export ----' )
    #symbolic_names = {0: "batch_size", 1: "max_seq_len"}
    symbolic_names = {0: "B", 1: "C", 2: "T"}
    torch.onnx.export(model,               # model being run
                      dummy_input,                         # model input (or a tuple for multiple inputs)
                      outModel,   # where to save the model (can be a file or file-like object)
                      export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=opsetVer,          # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      verbose=False,
                      input_names = ['signal'],   # the model's input names
                      output_names = ['segments'], # the model's output names
                      dynamic_axes={'signal' : symbolic_names,    # variable length axes
                                    })
    print( f'---- model has been saved to {outModel} ----\n' )

def verify():
    print( '\n\n\n====================verify=======================' )
    model =  Model.from_pretrained(
            "pyannote/segmentation@2022.07",
            strict=False,
            use_auth_token=hf_auth_token,
        )
    model.eval()

    onnx_model = onnx.load( outModel )
    onnx.checker.check_model(onnx_model)
    providers = ['CPUExecutionProvider']
    so = onnxruntime.SessionOptions()
    so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), so, providers=providers)

    wavs = [
            '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/multi-speaker_4-speakers_Jennifer_Aniston_and_Adam_Sandler_talk.wav',
            ]


    # Load wav
    sequence_len = 80000
    signal, fs = torchaudio.load( wavs[0] )
    assert( signal.shape[1] >= batch * sequence_len )
    signal = signal[:,:batch*sequence_len].flatten().reshape( batch, -1 )
    signal = signal.unsqueeze(1)

    # Calc using pytorch
    #with torch.no_grad():
    segments = model(signal)
    segments1 = segments.detach().numpy()

    # Calc embedding using onnx model
    print( ort_session.get_inputs()[0])
    ort_inputs = {ort_session.get_inputs()[0].name: signal.numpy()}
    ort_outs = ort_session.run(None, ort_inputs)
    onnx_segments = ort_outs[0]

    try:
        # Check if result is close enough
        np.testing.assert_allclose(
                segments1, 
                onnx_segments, 
                rtol=1e-03, 
                atol=1e-05,
                verbose = True)
        print( 'check passed' )
    except AssertionError as e:
        print( e )
        print( onnx_segments[0] )
    print( "-----------------------------" )

def debug():
    '''
    Load model
    check model
    print output shape
    print output 
    print outputs of all layers -- > this is very useful for debugging, 
        - check name against graph in netron.app 
        - python -m pdb export2.py to check each step input/output
    '''
    # Check exported model see if there is error
    # add all intermediate outputs to onnx net ------ 1
    ort_session = onnxruntime.InferenceSession(outModel)
    org_outputs = [x.name for x in ort_session.get_outputs()]

    print( f'\n---- Load model: {outModel} ----' )
    onnx_model = onnx.load( outModel )
    # Check model
    onnx.checker.check_model(onnx_model, full_check=True)
    inferred = onnx.shape_inference.infer_shapes(onnx_model, check_type=True)

    # Add all intermediate layer's output to model outputs
    for node in onnx_model.graph.node:
        for output in node.output:
            if output not in org_outputs:
                onnx_model.graph.output.extend([onnx.ValueInfoProto(name=output)])

    # Load test data
    signal, fs = torchaudio.load('/home/leo/storage/sharedFolderVirtualbox/audioForTesting/english_15s16k.wav')
    print( signal.shape )
    # batch size = 2
    signal = signal[:160000]

    model =  Model.from_pretrained(
            "pyannote/segmentation@2022.07",
            strict=False,
            use_auth_token=hf_auth_token,
        )
    model.eval()

    # Calc embeddings using pytorch
    #with torch.no_grad():
    segments = model(signal)


    # Compare the output of the original model and the ONNX-converted model to ensure their equivalence.
    providers = ['CPUExecutionProvider']
    so = onnxruntime.SessionOptions()
    #so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
    so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_BASIC
    so.log_severity_level = 1 
    ort_session = onnxruntime.InferenceSession(onnx_model.SerializeToString(), so, providers=providers)
    print( ort_session.get_inputs()[0])
    signal = signal[None,:]
    ort_inputs = {ort_session.get_inputs()[0].name: signal.numpy()}
    outputs = [x.name for x in ort_session.get_outputs()]
    #ort_inputs = {ort_session.get_inputs()[0].name: feats.numpy()}

    print( f'---- run model ----' )
    #ort_outs = ort_session.run(None, ort_inputs)
    ort_outs = ort_session.run(outputs, ort_inputs)
    onnx_segments = ort_outs[0]
    # ------ 3
    from collections import OrderedDict
    ort_outs = OrderedDict(zip(outputs, ort_outs))
    # Check output is same( or close enough )

    print( '---- print all layers outputs ---- ' )
    for name in outputs:
        print( f'--- out for {name} ---' )
        print( ort_outs[name] )

    print( f'---- check result ----' )
    #fn = torch.from_numpy( ort_outs[0] )
    #out = model.mods.mean_var_norm( fn, wav_lens)
    np.testing.assert_allclose(
            segments.detach().numpy(), 
            onnx_segments, 
            rtol=1e-03, 
            atol=1e-05,
            verbose = True)
    print( f'---- check result is done ----' )

    import sys
    np.set_printoptions(threshold=sys.maxsize)
    '''
    print("Original Output:")
    a1 = feats.detach().numpy().flatten()
    print("Onnx model Output:")
    a2 = ort_outs[0].flatten()
    if len( a1 ) != len( a2 ):
        print( f"Size diff: {len(a1)}, {len(a2)}" )
    for i in range( len( a1 )):
        if abs( a1[i] - a2[i] ) > 1e-03:
            print( f"Diff: {a1[i]}, {a2[i]}" )
            return
    '''


def hints():
    print( '\n\n\n====================note=======================' )
    print( 'If you error like below' )
    print('''
            RuntimeError: cuDNN version incompatibility: PyTorch was compiled  against (8, 7, 0) but found runtime version (8, 6, 0). PyTorch already comes bundled with cuDNN. One option to resolving this error is to ensure PyTorch can find the bundled cuDNN.Looks like your LD_LIBRARY_PATH contains incompatible version of cudnnPlease either remove it from the path or install cudnn (8, 7, 0)
            ''')
    print( 'then execute following command' )
    print( '$> unset LD_LIBRARY_PATH' )
    print( '===============================================\n\n\n' )


hints()
export()
#debug()
verify()

segment/export2.py
leohuang2013/pyannote-audio_speaker-diarization_cpp

# put your own hugging face token
hf_auth_token="YOUR_OWN_HF_TOKEN"

segment/hf_token.py
leohuang2013/pyannote-audio_speaker-diarization_cpp
from pyannote.audio.core.model import Model
from utils import SlidingWindow, SlidingWindowFeature, Segment
from typing import Optional, Tuple, Union
from hf_token import *
import torchaudio
import numpy as np
import torch
from einops import rearrange
import itertools
import math
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence
import torch.nn as nn

def batchify(iterable, batch_size: int = 32, fillvalue=None):
    """Batchify iterable"""
    # batchify('ABCDEFG', 3) --> ['A', 'B', 'C']  ['D', 'E', 'F']  [G, ]
    args = [iter(iterable)] * batch_size
    return itertools.zip_longest(*args, fillvalue=fillvalue)

class MySegment(nn.Module):
    def __init__(self, device:torch.device = None):
        super(MySegment, self).__init__()
        self.duration = 5.0
        self.step = 0.5
        self.batch_size = 32
        self.sample_rate = 16000
        self.embedding_batch_size = 32 # from  speaker_diarization config.yaml
        # minimum number of samples needed to extract an embedding
        # (a lower number of samples would result in an error)
        # This value determined by pyannote/audio/pipelines/speaker_verification.py:264
        self.min_num_samples = 640 # self._embedding.min_num_samples
        self.model =  Model.from_pretrained(
                "pyannote/segmentation@2022.07",
                strict=False,
                use_auth_token=hf_auth_token,
            )
        self.model.eval()
        self.model.to( torch.device( 'cuda' ))
        if device is None:
            self.device = self.model.device
        else:
            self.device = device

        # read from 
        # ~/.cache/torch/pyannote/models--pyannote--speaker-diarization/snapshots/2c6a571d14c3794623b098a065ff95fa22da7f25/config.yaml
        self.diarization_segmentation_threashold = 0.4442333667381752
        self.diarization_segmentation_min_duration_off = 0.5817029604921046

    #def slide(self, waveform: torch.Tensor, sample_rate: int) -> SlidingWindowFeature:
    def slide(self, waveform: torch.Tensor, sample_rate: int) -> torch.Tensor:
        """Slide model on a waveform

        Parameters
        ----------
        waveform: (num_channels, num_samples) torch.Tensor
            Waveform.
        sample_rate : int
            Sample rate.

        Returns
        -------
        output : SlidingWindowFeature
            Model output. Shape is (num_chunks, dimension) for chunk-level tasks,
            and (num_frames, dimension) for frame-level tasks.
        """

        window_size: int = round(self.duration * sample_rate) # 80000
        step_size: int = round(self.step * sample_rate) # 8000
        num_channels, num_samples = waveform.shape

        # Those code only to calc num_frames_per_chunk
        specifications = self.model.specifications
        introspection = self.model.introspection 
        num_frames_per_chunk, _ = introspection(window_size)

        # prepare complete chunks
        if num_samples >= window_size:
            chunks: torch.Tensor = rearrange(
                waveform.unfold(1, window_size, step_size),
                "channel chunk frame -> chunk channel frame",
            )
            num_chunks, _, _ = chunks.shape
        else:
            num_chunks = 0 

        # prepare last incomplete chunk
        has_last_chunk = (num_samples < window_size) or (
            num_samples - window_size
        ) % step_size > 0
        if has_last_chunk:
            last_chunk: torch.Tensor = waveform[:, num_chunks * step_size :]

        outputs: Union[List[np.ndarray], np.ndarray] = list()

        # slide over audio chunks in batch
        for c in np.arange(0, num_chunks, self.batch_size):
            batch: torch.Tensor = chunks[c : c + self.batch_size]
            outputs.append(self.infer(batch))

        # process orphan last chunk
        if has_last_chunk:
            # last_chunk 1x79674
            # last_output: numpy.array: 1x292x3
            last_output = self.infer(last_chunk[None])
            pad = num_frames_per_chunk - last_output.shape[1]
            last_output = np.pad(last_output, ((0, 0), (0, pad), (0, 0)))
            # after pad, last_output: numpy.array: 1x293x3

            outputs.append(last_output)
        # outputs in: list of numpy.array and first element is: 32x293x3
        # outputs out: list of numpy.array: 936x293x3
        outputs = np.vstack(outputs)

        # /home/leo/product/speaker_diarization/diarization/lib/python3.10/site-packages/pyannote/core/segment.py:441
        # /home/leo/product/speaker_diarization/diarization/lib/python3.10/site-packages/pyannote/core/feature.py:49
        frames = SlidingWindow(start=0.0, duration=self.duration, step=self.step)
        return SlidingWindowFeature(outputs, frames)

    def infer(self, chunks: torch.Tensor) -> np.ndarray:
        """Forward pass

        Takes care of sending chunks to right device and outputs back to CPU

        Parameters
        ----------
        chunks : (batch_size, num_channels, num_samples) torch.Tensor
            Batch of audio chunks.

        Returns
        -------
        outputs : (batch_size, ...) np.ndarray
            Model output.
        """

        with torch.no_grad():
            try:
                outputs = self.model(chunks.to(self.device))
            except RuntimeError as exception:
                if is_oom_error(exception):
                    raise MemoryError(
                        f"batch_size ({self.batch_size: d}) is probably too large. "
                        f"Try with a smaller value until memory error disappears."
                    )
                else:
                    raise exception

        return outputs.cpu().numpy()

    def embedding_mask(
        self, waveforms: torch.Tensor, masks: torch.Tensor = None
        ) -> Tuple[torch.Tensor, torch.Tensor]:
        """

        Parameters
        ----------
        waveforms : (batch_size, num_channels, num_samples)
            Only num_channels == 1 is supported.
        masks : (batch_size, num_samples), optional

        Returns
        -------
        embeddings : (batch_size, dimension)

        """

        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

            # TODO: speed up the creation of "signals"
            # preliminary profiling experiments show
            # that it accounts for 15% of __call__
            # (the remaining 85% being the actual forward pass)

            imasks = F.interpolate(
                masks.unsqueeze(dim=1), size=num_samples, mode="nearest"
            ).squeeze(dim=1)

            imasks = imasks > 0.5

            signals = pad_sequence(
                [waveform[imask] for waveform, imask in zip(waveforms, imasks)],
                batch_first=True,
            )
            wav_lens = imasks.sum(dim=1)

        max_len = wav_lens.max()

        # corner case: every signal is too short
        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

        return signals, wav_lens

        #**********************************************
        # This is the step call embedding model to get embeddings
        #**********************************************
        #embeddings = (
        #    self.classifier_.encode_batch(signals, wav_lens=wav_lens)
        #    .squeeze(dim=1)
        #    .cpu()
        #    .numpy()
        #)

        #embeddings[too_short.cpu().numpy()] = np.NAN

        #return embeddings


    # from pyannote/audio/core/io.py
    def crop(
        self,
        waveform: torch.Tensor,
        sample_rate: int,
        segment: Segment,
        duration: Optional[float] = None,
        mode="raise",
    ) -> Tuple[torch.Tensor, int]:

        # infer which samples to load from sample rate and requested chunk
        start_frame = math.floor(segment.start * sample_rate)
        frames = waveform.shape[1]

        if duration:
            num_frames = math.floor(duration * sample_rate)
            end_frame = start_frame + num_frames
        else:
            end_frame = math.floor(segment.end * sample_rate)
            num_frames = end_frame - start_frame

        pad_start = -min(0, start_frame)
        pad_end = max(end_frame, frames) - frames
        start_frame = max(0, start_frame)
        end_frame = min(end_frame, frames)
        num_frames = end_frame - start_frame

        data = waveform[:, start_frame:end_frame]

        # pad with zeros
        if mode == "pad":
            data = F.pad(data, (pad_start, pad_end))

        return self.downmix_and_resample(data, sample_rate)

    # from pyannote/audio/core/io.py
    def downmix_and_resample(self, waveform: torch.Tensor, sample_rate: int) -> torch.Tensor:
        """Downmix and resample

        Parameters
        ----------
        waveform : (channel, time) Tensor
            Waveform.
        sample_rate : int
            Sample rate.

        Returns
        -------
        waveform : (channel, time) Tensor
            Remixed and resampled waveform
        sample_rate : int
            New sample rate
        """

        # downmix to mono
        #if self.mono and waveform.shape[0] > 1:
        if waveform.shape[0] > 1:
            waveform = waveform.mean(dim=0, keepdim=True)

        # resample
        if (self.sample_rate is not None) and (self.sample_rate != sample_rate):
            waveform = torchaudio.functional.resample(
                waveform, sample_rate, self.sample_rate
            )
            sample_rate = self.sample_rate

        return waveform, sample_rate


    def binarize_swf( 
            self,
        scores: SlidingWindowFeature,
        onset: float = 0.5,
        offset: Optional[float] = None,
        initial_state: Optional[bool] = None,
    ):
        """(Batch) hysteresis thresholding

        Parameters
        ----------
        scores : SlidingWindowFeature
            (num_chunks, num_frames, num_classes)- or (num_frames, num_classes)-shaped scores.
        onset : float, optional
            Onset threshold. Defaults to 0.5.
        offset : float, optional
            Offset threshold. Defaults to `onset`.
        initial_state : np.ndarray or bool, optional
            Initial state.

        Returns
        -------
        binarized : same as scores
            Binarized scores with same shape and type as scores.

        """

        offset = offset or onset

        if scores.data.ndim == 2:
            num_frames, num_classes = scores.data.shape
            data = einops.rearrange(scores.data, "f k -> k f", f=num_frames, k=num_classes)
            binarized = self.binarize_ndarray(
                data, onset=onset, offset=offset, initial_state=initial_state
            )
            return SlidingWindowFeature(
                1.0
                * einops.rearrange(binarized, "k f -> f k", f=num_frames, k=num_classes),
                scores.sliding_window,
            )

        elif scores.data.ndim == 3:
            num_chunks, num_frames, num_classes = scores.data.shape
            data = rearrange(
                scores.data, "c f k -> (c k) f", c=num_chunks, f=num_frames, k=num_classes
            )
            binarized = self.binarize_ndarray(
                data, onset=onset, offset=offset, initial_state=initial_state
            )
            return SlidingWindowFeature(
                1.0
                * rearrange(
                    binarized, "(c k) f -> c f k", c=num_chunks, f=num_frames, k=num_classes
                ),
                scores.sliding_window,
            )

        else:
            raise ValueError(
                "Shape of scores must be (num_chunks, num_frames, num_classes) or (num_frames, num_classes)."
            )

    def binarize_ndarray(
            self,
            scores: np.ndarray,
            onset: float = 0.5,
            offset: Optional[float] = None,
            initial_state: Optional[Union[bool, np.ndarray]] = None,
        ):
        """(Batch) hysteresis thresholding

        Parameters
        ----------
        scores : numpy.ndarray
            (num_frames, num_classes)-shaped scores.
        onset : float, optional
            Onset threshold. Defaults to 0.5.
        offset : float, optional
            Offset threshold. Defaults to `onset`.
        initial_state : np.ndarray or bool, optional
            Initial state.

        Returns
        -------
        binarized : same as scores
            Binarized scores with same shape and type as scores.
        """

        offset = offset or onset

        batch_size, num_frames = scores.shape

        scores = np.nan_to_num(scores)

        if initial_state is None:
            initial_state = scores[:, 0] >= 0.5 * (onset + offset)

        elif isinstance(initial_state, bool):
            initial_state = initial_state * np.ones((batch_size,), dtype=bool)

        elif isinstance(initial_state, np.ndarray):
            assert initial_state.shape == (batch_size,)
            assert initial_state.dtype == bool

        initial_state = np.tile(initial_state, (num_frames, 1)).T

        on = scores > onset
        off_or_on = (scores < offset) | on

        # indices of frames for which the on/off state is well-defined
        well_defined_idx = np.array(
            list(itertools.zip_longest(*[np.nonzero(oon)[0] for oon in off_or_on], fillvalue=-1))
        ).T

        # corner case where well_defined_idx is empty
        if not well_defined_idx.size:
            return np.zeros_like(scores, dtype=bool) | initial_state

        # points to the index of the previous well-defined frame
        same_as = np.cumsum(off_or_on, axis=1)

        samples = np.tile(np.arange(batch_size), (num_frames, 1)).T

        return np.where(
            same_as, on[samples, well_defined_idx[samples, same_as - 1]], initial_state
        )

    def forward( self, waveform:torch.Tensor ):
        # waveform, sample_rate = self.model.audio(file)
        sample_rate = self.sample_rate
        segmentations = self.slide( waveform, sample_rate ) 
        return segmentations # for testing
        #   shape: (num_chunks, num_frames, local_num_speakers)

        # estimate frame-level number of instantaneous speakers
        #count = self.speaker_count(
        #    segmentations,
        #    onset=self.diarization_segmentation_threashold,
        #    frames=self._frames,
        #)
        #   shape: (num_frames, 1)
        #   dtype: int

        # binarize segmentation
        binary_segmentations: SlidingWindowFeature = self.binarize_swf(
            segmentations,
            onset=self.diarization_segmentation_threashold,
            initial_state=False,
        )

        # Next, pass binarized_segmentations and audio file to embedding module
        # pyannote/audio/pipelines/speaker_diarization.py: get_embeddings(...)
        duration = binary_segmentations.sliding_window.duration
        num_chunks, num_frames, _ = binary_segmentations.data.shape

        # corresponding minimum number of frames
        num_samples = duration * 16000 # self._embedding.sample_rate
        min_num_frames = math.ceil(num_frames * self.min_num_samples / num_samples)

        # zero-out frames with overlapping speech
        clean_frames = 1.0 * (
            np.sum(binary_segmentations.data, axis=2, keepdims=True) < 2
        )
        clean_segmentations = SlidingWindowFeature(
            binary_segmentations.data * clean_frames,
            binary_segmentations.sliding_window,
        )

        def iter_waveform_and_mask():
            for (chunk, masks), (_, clean_masks) in zip(
                binary_segmentations, clean_segmentations
            ):
                # chunk: Segment(t, t + duration)
                # masks: (num_frames, local_num_speakers) np.ndarray

                wvform, _ = self.crop(
                    waveform,
                    sample_rate,
                    chunk,
                    duration=duration,
                    mode="pad",
                )
                # wvform: (1, num_samples) torch.Tensor

                # mask may contain NaN (in case of partial stitching)
                masks = np.nan_to_num(masks, nan=0.0).astype(np.float32)
                clean_masks = np.nan_to_num(clean_masks, nan=0.0).astype(np.float32)

                for mask, clean_mask in zip(masks.T, clean_masks.T):
                    # mask: (num_frames, ) np.ndarray

                    if np.sum(clean_mask) > min_num_frames:
                        used_mask = clean_mask
                    else:
                        used_mask = mask

                    yield wvform[None], torch.from_numpy(used_mask)[None]
                    # w: (1, 1, num_samples) torch.Tensor
                    # m: (1, num_frames) torch.Tensor

        batches = batchify(
            iter_waveform_and_mask(),
            batch_size=self.embedding_batch_size,
            fillvalue=(None, None),
        )

        masked_waveforms_batches = []

        for batch in batches:
            waveforms, masks = zip(*filter(lambda b: b[0] is not None, batch))

            # waveforms is tuple of tensors which is 1x1x80000
            #waveform_batch = torch.vstack(waveforms) --> changed to below
            # waveform_batch is tensor: 32x1x80000
            # (batch_size, 1, num_samples) torch.Tensor
            tmp = []
            for wf in waveforms:
                tmp.append( wf.numpy())
            waveform_batch = torch.from_numpy( np.vstack( tmp ))

            #mask_batch = torch.vstack(masks) --> changed to below
            # (batch_size, num_frames) torch.Tensor
            tmp = []
            for wf in masks:
                tmp.append( wf.numpy())
            mask_batch = torch.from_numpy( np.vstack( tmp ))

            masked_waveforms, wav_len = self.embedding_mask(
                waveform_batch, masks = mask_batch
            )
            # (batch_size, dimension) np.ndarray
            #masked_waveforms_batches.append( masked_waveforms.flatten() )
            masked_waveforms_batches.append( masked_waveforms )

        #xnew = torch.stack((masked_waveforms_batches[0], masked_waveforms_batches[0]), 0)

        return masked_waveforms_batches
        #return torch.from_numpy( np.array( aa ))
        #return xnew



segment/mysegment.py
leohuang2013/pyannote-audio_speaker-diarization_cpp
from pyannote.audio.core.model import Model
from utils import SlidingWindow, SlidingWindowFeature, Segment
from typing import Optional, Tuple, Union
from hf_token import *
import torchaudio
import numpy as np
import torch
from einops import rearrange
import itertools
import math
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence

def batchify(iterable, batch_size: int = 32, fillvalue=None):
    """Batchify iterable"""
    # batchify('ABCDEFG', 3) --> ['A', 'B', 'C']  ['D', 'E', 'F']  [G, ]
    args = [iter(iterable)] * batch_size
    return itertools.zip_longest(*args, fillvalue=fillvalue)


class Segmentation:
    def __init__(self, device:torch.device = None):
        self.duration = 5.0
        self.step = 0.5
        self.batch_size = 32
        self.sample_rate = 16000
        self.embedding_batch_size = 32 # from  speaker_diarization config.yaml
        # minimum number of samples needed to extract an embedding
        # (a lower number of samples would result in an error)
        # This value determined by pyannote/audio/pipelines/speaker_verification.py:264
        self.min_num_samples = 640 # self._embedding.min_num_samples
        self.model =  Model.from_pretrained(
                "pyannote/segmentation@2022.07",
                strict=False,
                use_auth_token=hf_auth_token,
            )
        self.model.eval()
        self.model.to( torch.device( 'cuda' ))
        if device is None:
            self.device = self.model.device
        else:
            self.device = device
        #self.device = torch.device('cuda')

        # read from 
        # ~/.cache/torch/pyannote/models--pyannote--speaker-diarization/snapshots/2c6a571d14c3794623b098a065ff95fa22da7f25/config.yaml
        self.diarization_segmentation_threashold = 0.4442333667381752
        self.diarization_segmentation_min_duration_off = 0.5817029604921046



    def slide(self, waveform: torch.Tensor, sample_rate: int) -> SlidingWindowFeature:
        """Slide model on a waveform

        Parameters
        ----------
        waveform: (num_channels, num_samples) torch.Tensor
            Waveform.
        sample_rate : int
            Sample rate.

        Returns
        -------
        output : SlidingWindowFeature
            Model output. Shape is (num_chunks, dimension) for chunk-level tasks,
            and (num_frames, dimension) for frame-level tasks.
        """

        window_size: int = round(self.duration * sample_rate)
        step_size: int = round(self.step * sample_rate)
        num_channels, num_samples = waveform.shape

        # Those code only to calc num_frames_per_chunk
        specifications = self.model.specifications
        introspection = self.model.introspection 
        num_frames_per_chunk, _ = introspection(window_size)

        # prepare complete chunks
        if num_samples >= window_size:
            chunks: torch.Tensor = rearrange(
                waveform.unfold(1, window_size, step_size),
                "channel chunk frame -> chunk channel frame",
            )
            num_chunks, _, _ = chunks.shape
        else:
            num_chunks = 0 

        # prepare last incomplete chunk
        has_last_chunk = (num_samples < window_size) or (
            num_samples - window_size
        ) % step_size > 0
        if has_last_chunk:
            last_chunk: torch.Tensor = waveform[:, num_chunks * step_size :]

        outputs: Union[List[np.ndarray], np.ndarray] = list()

        # slide over audio chunks in batch
        for c in np.arange(0, num_chunks, self.batch_size):
            batch: torch.Tensor = chunks[c : c + self.batch_size]
            outputs.append(self.infer(batch))

        # process orphan last chunk
        if has_last_chunk:
            last_output = self.infer(last_chunk[None])
            pad = num_frames_per_chunk - last_output.shape[1]
            last_output = np.pad(last_output, ((0, 0), (0, pad), (0, 0)))

            outputs.append(last_output)

        outputs = np.vstack(outputs)

        # /home/leo/product/speaker_diarization/diarization/lib/python3.10/site-packages/pyannote/core/segment.py:441
        # /home/leo/product/speaker_diarization/diarization/lib/python3.10/site-packages/pyannote/core/feature.py:49
        frames = SlidingWindow(start=0.0, duration=self.duration, step=self.step)
        return SlidingWindowFeature(outputs, frames)

    def infer(self, chunks: torch.Tensor) -> np.ndarray:
        """Forward pass

        Takes care of sending chunks to right device and outputs back to CPU

        Parameters
        ----------
        chunks : (batch_size, num_channels, num_samples) torch.Tensor
            Batch of audio chunks.

        Returns
        -------
        outputs : (batch_size, ...) np.ndarray
            Model output.
        """

        with torch.no_grad():
            try:
                outputs = self.model(chunks.to(self.device))
            except RuntimeError as exception:
                if is_oom_error(exception):
                    raise MemoryError(
                        f"batch_size ({self.batch_size: d}) is probably too large. "
                        f"Try with a smaller value until memory error disappears."
                    )
                else:
                    raise exception

        return outputs.cpu().numpy()


    def run( self, waveform:torch.Tensor, sample_rate:int ):
        # waveform, sample_rate = self.model.audio(file)
        segmentations = self.slide( waveform, sample_rate ) 
        #   shape: (num_chunks, num_frames, local_num_speakers)

        # estimate frame-level number of instantaneous speakers
        #count = self.speaker_count(
        #    segmentations,
        #    onset=self.diarization_segmentation_threashold,
        #    frames=self._frames,
        #)
        #   shape: (num_frames, 1)
        #   dtype: int

        # binarize segmentation
        binary_segmentations: SlidingWindowFeature = self.binarize_swf(
            segmentations,
            onset=self.diarization_segmentation_threashold,
            initial_state=False,
        )

        # Next, pass binarized_segmentations and audio file to embedding module
        # pyannote/audio/pipelines/speaker_diarization.py: get_embeddings(...)
        duration = binary_segmentations.sliding_window.duration
        num_chunks, num_frames, _ = binary_segmentations.data.shape

        # corresponding minimum number of frames
        num_samples = duration * 16000 # self._embedding.sample_rate
        min_num_frames = math.ceil(num_frames * self.min_num_samples / num_samples)

        # zero-out frames with overlapping speech
        clean_frames = 1.0 * (
            np.sum(binary_segmentations.data, axis=2, keepdims=True) < 2
        )
        clean_segmentations = SlidingWindowFeature(
            binary_segmentations.data * clean_frames,
            binary_segmentations.sliding_window,
        )

        def iter_waveform_and_mask():
            for (chunk, masks), (_, clean_masks) in zip(
                binary_segmentations, clean_segmentations
            ):
                # chunk: Segment(t, t + duration)
                # masks: (num_frames, local_num_speakers) np.ndarray

                wvform, _ = self.crop(
                    waveform,
                    sample_rate,
                    chunk,
                    duration=duration,
                    mode="pad",
                )
                # wvform: (1, num_samples) torch.Tensor

                # mask may contain NaN (in case of partial stitching)
                masks = np.nan_to_num(masks, nan=0.0).astype(np.float32)
                clean_masks = np.nan_to_num(clean_masks, nan=0.0).astype(np.float32)

                for mask, clean_mask in zip(masks.T, clean_masks.T):
                    # mask: (num_frames, ) np.ndarray

                    if np.sum(clean_mask) > min_num_frames:
                        used_mask = clean_mask
                    else:
                        used_mask = mask

                    yield wvform[None], torch.from_numpy(used_mask)[None]
                    # w: (1, 1, num_samples) torch.Tensor
                    # m: (1, num_frames) torch.Tensor

        batches = batchify(
            iter_waveform_and_mask(),
            batch_size=self.embedding_batch_size,
            fillvalue=(None, None),
        )

        embedding_batches = []

        for batch in batches:
            waveforms, masks = zip(*filter(lambda b: b[0] is not None, batch))

            waveform_batch = torch.vstack(waveforms)
            # (batch_size, 1, num_samples) torch.Tensor

            mask_batch = torch.vstack(masks)
            # (batch_size, num_frames) torch.Tensor

            embedding_batch: np.ndarray = self.embedding_mask(
                waveform_batch, masks=mask_batch
            )
            # (batch_size, dimension) np.ndarray
            embedding_batches.append(embedding_batch)

        #embedding_batches = np.vstack(embedding_batches)
        #embeddings = rearrange(embedding_batches, "(c s) d -> c s d", c=num_chunks)

        #return embeddings

    def embedding_mask(
        self, waveforms: torch.Tensor, masks: torch.Tensor = None
        ) -> np.ndarray:
        """

        Parameters
        ----------
        waveforms : (batch_size, num_channels, num_samples)
            Only num_channels == 1 is supported.
        masks : (batch_size, num_samples), optional

        Returns
        -------
        embeddings : (batch_size, dimension)

        """

        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

            # TODO: speed up the creation of "signals"
            # preliminary profiling experiments show
            # that it accounts for 15% of __call__
            # (the remaining 85% being the actual forward pass)

            imasks = F.interpolate(
                masks.unsqueeze(dim=1), size=num_samples, mode="nearest"
            ).squeeze(dim=1)

            imasks = imasks > 0.5

            signals = pad_sequence(
                [waveform[imask] for waveform, imask in zip(waveforms, imasks)],
                batch_first=True,
            )
            wav_lens = imasks.sum(dim=1)

        max_len = wav_lens.max()

        # corner case: every signal is too short
        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

        #**********************************************
        # This is the step call embedding model to get embeddings
        #**********************************************
        #embeddings = (
        #    self.classifier_.encode_batch(signals, wav_lens=wav_lens)
        #    .squeeze(dim=1)
        #    .cpu()
        #    .numpy()
        #)

        #embeddings[too_short.cpu().numpy()] = np.NAN

        #return embeddings
        return None


    # from pyannote/audio/core/io.py
    def crop(
        self,
        waveform: torch.Tensor,
        sample_rate: int,
        segment: Segment,
        duration: Optional[float] = None,
        mode="raise",
    ) -> Tuple[torch.Tensor, int]:

        # infer which samples to load from sample rate and requested chunk
        start_frame = math.floor(segment.start * sample_rate)
        frames = waveform.shape[1]

        if duration:
            num_frames = math.floor(duration * sample_rate)
            end_frame = start_frame + num_frames
        else:
            end_frame = math.floor(segment.end * sample_rate)
            num_frames = end_frame - start_frame

        pad_start = -min(0, start_frame)
        pad_end = max(end_frame, frames) - frames
        start_frame = max(0, start_frame)
        end_frame = min(end_frame, frames)
        num_frames = end_frame - start_frame

        data = waveform[:, start_frame:end_frame]

        # pad with zeros
        if mode == "pad":
            data = F.pad(data, (pad_start, pad_end))

        return data, sample_rate


    def binarize_swf( 
            self,
        scores: SlidingWindowFeature,
        onset: float = 0.5,
        offset: Optional[float] = None,
        initial_state: Optional[bool] = None,
    ):
        """(Batch) hysteresis thresholding

        Parameters
        ----------
        scores : SlidingWindowFeature
            (num_chunks, num_frames, num_classes)- or (num_frames, num_classes)-shaped scores.
        onset : float, optional
            Onset threshold. Defaults to 0.5.
        offset : float, optional
            Offset threshold. Defaults to `onset`.
        initial_state : np.ndarray or bool, optional
            Initial state.

        Returns
        -------
        binarized : same as scores
            Binarized scores with same shape and type as scores.

        """

        offset = offset or onset

        if scores.data.ndim == 2:
            num_frames, num_classes = scores.data.shape
            data = einops.rearrange(scores.data, "f k -> k f", f=num_frames, k=num_classes)
            binarized = self.binarize_ndarray(
                data, onset=onset, offset=offset, initial_state=initial_state
            )
            return SlidingWindowFeature(
                1.0
                * einops.rearrange(binarized, "k f -> f k", f=num_frames, k=num_classes),
                scores.sliding_window,
            )

        elif scores.data.ndim == 3:
            num_chunks, num_frames, num_classes = scores.data.shape
            data = rearrange(
                scores.data, "c f k -> (c k) f", c=num_chunks, f=num_frames, k=num_classes
            )
            binarized = self.binarize_ndarray(
                data, onset=onset, offset=offset, initial_state=initial_state
            )
            return SlidingWindowFeature(
                1.0
                * rearrange(
                    binarized, "(c k) f -> c f k", c=num_chunks, f=num_frames, k=num_classes
                ),
                scores.sliding_window,
            )

        else:
            raise ValueError(
                "Shape of scores must be (num_chunks, num_frames, num_classes) or (num_frames, num_classes)."
            )

    def binarize_ndarray(
            self,
            scores: np.ndarray,
            onset: float = 0.5,
            offset: Optional[float] = None,
            initial_state: Optional[Union[bool, np.ndarray]] = None,
        ):
        """(Batch) hysteresis thresholding

        Parameters
        ----------
        scores : numpy.ndarray
            (num_frames, num_classes)-shaped scores.
        onset : float, optional
            Onset threshold. Defaults to 0.5.
        offset : float, optional
            Offset threshold. Defaults to `onset`.
        initial_state : np.ndarray or bool, optional
            Initial state.

        Returns
        -------
        binarized : same as scores
            Binarized scores with same shape and type as scores.
        """

        offset = offset or onset

        batch_size, num_frames = scores.shape

        scores = np.nan_to_num(scores)

        if initial_state is None:
            initial_state = scores[:, 0] >= 0.5 * (onset + offset)

        elif isinstance(initial_state, bool):
            initial_state = initial_state * np.ones((batch_size,), dtype=bool)

        elif isinstance(initial_state, np.ndarray):
            assert initial_state.shape == (batch_size,)
            assert initial_state.dtype == bool

        initial_state = np.tile(initial_state, (num_frames, 1)).T

        on = scores > onset
        off_or_on = (scores < offset) | on

        # indices of frames for which the on/off state is well-defined
        well_defined_idx = np.array(
            list(itertools.zip_longest(*[np.nonzero(oon)[0] for oon in off_or_on], fillvalue=-1))
        ).T

        # corner case where well_defined_idx is empty
        if not well_defined_idx.size:
            return np.zeros_like(scores, dtype=bool) | initial_state

        # points to the index of the previous well-defined frame
        same_as = np.cumsum(off_or_on, axis=1)

        samples = np.tile(np.arange(batch_size), (num_frames, 1)).T

        return np.where(
            same_as, on[samples, well_defined_idx[samples, same_as - 1]], initial_state
        )




if __name__ == '__main__':
    #audio = '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/one_min16k.wav')
    #audio = '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/multi-speaker_7-speakers_SkillsFuture.wav'
    audio = '/home/leo/storage/sharedFolderVirtualbox/audioForTesting/multi-speaker_4-speakers_Jennifer_Aniston_and_Adam_Sandler_talk.wav'
    signal, fs = torchaudio.load( audio )
    seg = Segmentation()
    seg.run( signal, fs )

segment/segment.py
leohuang2013/pyannote-audio_speaker-diarization_cpp
import torch.nn as nn
import torch
from typing import Tuple
import numpy as np

class MySimple(nn.Module):
    def __init__(self, device:torch.device = None):
        super(MySimple, self).__init__()

    def calc( self, waveforms:torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]:
        r1 = waveforms.squeeze(dim=1)
        r2 = waveforms.T
        return r1, r2

    def slide(self, waveform: torch.Tensor, sample_rate: int) -> torch.Tensor:
        ws = 16
        step = 4
        d = waveform.unfold( 1, ws, step )
        #i = 0
        #a = []
        #while i + ws < waveform.shape[1]:
        #    a.append( waveform[0][i:i+ws] )
        #    i += step
        #d = torch.stack( tuple( a ))
        #c = np.reshape( n, ( 32, 1, -1 ))
        #d = torch.from_numpy( n ) // failed
        #d = torch.reshape( waveform, ( sample_rate, 1, -1 )) // working
        return d


    def forward( self, waveform:torch.Tensor ):
        res = []
        r1, r2 = self.calc( waveform )
        res.append( r1 )
        res.append( r2 )
        feat = self.slide( waveform, 32 )
        return feat

opsetVer = 17
outModel = 'simple.onnx'

def export():
    # Create dummy input
    signal = torch.randn( 1, 128 ) 
    print( signal.shape )

    model = MySimple()

    # Export the model
    print( '\n---- start export ----' )
    symbolic_names = {0: "batch_size", 1: "max_seq_len"}
    torch.onnx.export(model,               # model being run
                      signal,                         # model input (or a tuple for multiple inputs)
                      outModel,   # where to save the model (can be a file or file-like object)
                      export_params=True,        # store the trained parameter weights inside the model file
                      opset_version=opsetVer,          # the ONNX version to export the model to
                      do_constant_folding=True,  # whether to execute constant folding for optimization
                      verbose=False,
                      input_names = ['signal'],   # the model's input names
                      output_names = ['segments'], # the model's output names
                      dynamic_axes={'signal' : symbolic_names,    # variable length axes
                                    })
    print( f'---- model has been saved to {outModel} ----\n' )


export()

segment/simple.py
leohuang2013/pyannote-audio_speaker-diarization_cpp

import numbers
import warnings
from typing import Tuple, Optional, Union, Iterator, Iterable, List, Text
#from typing import Dict, Tuple, Iterable, List, TYPE_CHECKING
from dataclasses import dataclass
import numpy as np

# setting 'frozen' to True makes it hashable and immutable
@dataclass(frozen=True, order=True)
class Segment:
    """
    Time interval

    Parameters
    ----------
    start : float
        interval start time, in seconds.
    end : float
        interval end time, in seconds.


    Segments can be compared and sorted using the standard operators:

    >>> Segment(0, 1) == Segment(0, 1.)
    True
    >>> Segment(0, 1) != Segment(3, 4)
    True
    >>> Segment(0, 1) < Segment(2, 3)
    True
    >>> Segment(0, 1) < Segment(0, 2)
    True
    >>> Segment(1, 2) < Segment(0, 3)
    False

    Note
    ----
    A segment is smaller than another segment if one of these two conditions is verified:

      - `segment.start < other_segment.start`
      - `segment.start == other_segment.start` and `segment.end < other_segment.end`

    """
    start: float = 0.0
    end: float = 0.0

    @staticmethod
    def set_precision(ndigits: Optional[int] = None):
        """Automatically round start and end timestamps to `ndigits` precision after the decimal point

        To ensure consistency between `Segment` instances, it is recommended to call this method only 
        once, right after importing `pyannote.core.Segment`.

        Usage
        -----
        >>> from pyannote.core import Segment
        >>> Segment.set_precision(2)
        >>> Segment(1/3, 2/3)
        <Segment(0.33, 0.67)>
        """
        global AUTO_ROUND_TIME
        global SEGMENT_PRECISION

        if ndigits is None:
            # backward compatibility
            AUTO_ROUND_TIME = False
            # 1 μs (one microsecond)
            SEGMENT_PRECISION = 1e-6
        else:
            AUTO_ROUND_TIME = True
            SEGMENT_PRECISION = 10 ** (-ndigits)

    def __bool__(self):
        """Emptiness

        >>> if segment:
        ...    # segment is not empty.
        ... else:
        ...    # segment is empty.

        Note
        ----
        A segment is considered empty if its end time is smaller than its
        start time, or its duration is smaller than 1μs.
        """
        return bool((self.end - self.start) > SEGMENT_PRECISION)

    #def __post_init__(self):
    #    """Round start and end up to SEGMENT_PRECISION precision (when required)"""
    #    if AUTO_ROUND_TIME:
    #        object.__setattr__(self, 'start', int(self.start / SEGMENT_PRECISION + 0.5) * SEGMENT_PRECISION)
    #        object.__setattr__(self, 'end', int(self.end / SEGMENT_PRECISION + 0.5) * SEGMENT_PRECISION)

    @property
    def duration(self) -> float:
        """Segment duration (read-only)"""
        return self.end - self.start if self else 0.

    @property
    def middle(self) -> float:
        """Segment mid-time (read-only)"""
        return .5 * (self.start + self.end)

    def __iter__(self) -> Iterator[float]:
        """Unpack segment boundaries
        >>> segment = Segment(start, end)
        >>> start, end = segment
        """
        yield self.start
        yield self.end

    def copy(self) -> 'Segment':
        """Get a copy of the segment

        Returns
        -------
        copy : Segment
            Copy of the segment.
        """
        return Segment(start=self.start, end=self.end)

    # ------------------------------------------------------- #
    # Inclusion (in), intersection (&), union (|) and gap (^) #
    # ------------------------------------------------------- #

    def __contains__(self, other: 'Segment'):
        """Inclusion

        >>> segment = Segment(start=0, end=10)
        >>> Segment(start=3, end=10) in segment:
        True
        >>> Segment(start=5, end=15) in segment:
        False
        """
        return (self.start <= other.start) and (self.end >= other.end)

    def __and__(self, other):
        """Intersection

        >>> segment = Segment(0, 10)
        >>> other_segment = Segment(5, 15)
        >>> segment & other_segment
        <Segment(5, 10)>

        Note
        ----
        When the intersection is empty, an empty segment is returned:

        >>> segment = Segment(0, 10)
        >>> other_segment = Segment(15, 20)
        >>> intersection = segment & other_segment
        >>> if not intersection:
        ...    # intersection is empty.
        """
        start = max(self.start, other.start)
        end = min(self.end, other.end)
        return Segment(start=start, end=end)

    def intersects(self, other: 'Segment') -> bool:
        """Check whether two segments intersect each other

        Parameters
        ----------
        other : Segment
            Other segment

        Returns
        -------
        intersect : bool
            True if segments intersect, False otherwise
        """

        return (self.start < other.start and
                other.start < self.end - SEGMENT_PRECISION) or \
               (self.start > other.start and
                self.start < other.end - SEGMENT_PRECISION) or \
               (self.start == other.start)

    def overlaps(self, t: float) -> bool:
        """Check if segment overlaps a given time

        Parameters
        ----------
        t : float
            Time, in seconds.

        Returns
        -------
        overlap: bool
            True if segment overlaps time t, False otherwise.
        """
        return self.start <= t and self.end >= t

    def __or__(self, other: 'Segment') -> 'Segment':
        """Union

        >>> segment = Segment(0, 10)
        >>> other_segment = Segment(5, 15)
        >>> segment | other_segment
        <Segment(0, 15)>

        Note
        ----
        When a gap exists between the segment, their union covers the gap as well:

        >>> segment = Segment(0, 10)
        >>> other_segment = Segment(15, 20)
        >>> segment | other_segment
        <Segment(0, 20)
        """

        # if segment is empty, union is the other one
        if not self:
            return other
        # if other one is empty, union is self
        if not other:
            return self

        # otherwise, do what's meant to be...
        start = min(self.start, other.start)
        end = max(self.end, other.end)
        return Segment(start=start, end=end)

    def __xor__(self, other: 'Segment') -> 'Segment':
        """Gap

        >>> segment = Segment(0, 10)
        >>> other_segment = Segment(15, 20)
        >>> segment ^ other_segment
        <Segment(10, 15)

        Note
        ----
        The gap between a segment and an empty segment is not defined.

        >>> segment = Segment(0, 10)
        >>> empty_segment = Segment(11, 11)
        >>> segment ^ empty_segment
        ValueError: The gap between a segment and an empty segment is not defined.
        """

        # if segment is empty, xor is not defined
        if (not self) or (not other):
            raise ValueError(
                'The gap between a segment and an empty segment '
                'is not defined.')

        start = min(self.end, other.end)
        end = max(self.start, other.start)
        return Segment(start=start, end=end)

    def _str_helper(self, seconds: float) -> str:
        from datetime import timedelta
        negative = seconds < 0
        seconds = abs(seconds)
        td = timedelta(seconds=seconds)
        seconds = td.seconds + 86400 * td.days
        microseconds = td.microseconds
        hours, remainder = divmod(seconds, 3600)
        minutes, seconds = divmod(remainder, 60)
        return '%s%02d:%02d:%02d.%03d' % (
            '-' if negative else ' ', hours, minutes,
            seconds, microseconds / 1000)

    def __str__(self):
        """Human-readable representation

        >>> print(Segment(1337, 1337 + 0.42))
        [ 00:22:17.000 -->  00:22:17.420]

        Note
        ----
        Empty segments are printed as "[]"
        """
        if self:
            return '[%s --> %s]' % (self._str_helper(self.start),
                                    self._str_helper(self.end))
        return '[]'

    def __repr__(self):
        """Computer-readable representation

        >>> Segment(1337, 1337 + 0.42)
        <Segment(1337, 1337.42)>
        """
        return '<Segment(%g, %g)>' % (self.start, self.end)

    def for_json(self):
        """Serialization

        See also
        --------
        :mod:`pyannote.core.json`
        """
        return {'start': self.start, 'end': self.end}

    @classmethod
    def from_json(cls, data):
        """Deserialization

        See also
        --------
        :mod:`pyannote.core.json`
        """
        return cls(start=data['start'], end=data['end'])

    def _repr_png_(self):
        """IPython notebook support

        See also
        --------
        :mod:`pyannote.core.notebook`
        """
        from .notebook import MATPLOTLIB_IS_AVAILABLE, MATPLOTLIB_WARNING
        if not MATPLOTLIB_IS_AVAILABLE:
            warnings.warn(MATPLOTLIB_WARNING.format(klass=self.__class__.__name__))
            return None

        from .notebook import repr_segment
        try:
            return repr_segment(self)
        except ImportError:
            warnings.warn(
                f"Couldn't import matplotlib to render the vizualization for object {self}. To enable, install the required dependencies with 'pip install pyannore.core[notebook]'")
            return None



class SlidingWindow:
    """Sliding window

    Parameters
    ----------
    duration : float > 0, optional
        Window duration, in seconds. Default is 30 ms.
    step : float > 0, optional
        Step between two consecutive position, in seconds. Default is 10 ms.
    start : float, optional
        First start position of window, in seconds. Default is 0.
    end : float > `start`, optional
        Default is infinity (ie. window keeps sliding forever)

    Examples
    --------

    >>> sw = SlidingWindow(duration, step, start)
    >>> frame_range = (a, b)
    >>> frame_range == sw.toFrameRange(sw.toSegment(*frame_range))
    ... True

    >>> segment = Segment(A, B)
    >>> new_segment = sw.toSegment(*sw.toFrameRange(segment))
    >>> abs(segment) - abs(segment & new_segment) < .5 * sw.step

    >>> sw = SlidingWindow(end=0.1)
    >>> print(next(sw))
    [ 00:00:00.000 -->  00:00:00.030]
    >>> print(next(sw))
    [ 00:00:00.010 -->  00:00:00.040]
    """

    def __init__(self, duration=0.030, step=0.010, start=0.000, end=None):

        # duration must be a float > 0
        if duration <= 0:
            raise ValueError("'duration' must be a float > 0.")
        self.__duration = duration

        # step must be a float > 0
        if step <= 0:
            raise ValueError("'step' must be a float > 0.")
        self.__step: float = step

        # start must be a float.
        self.__start: float = start

        # if end is not provided, set it to infinity
        if end is None:
            self.__end: float = np.inf
        else:
            # end must be greater than start
            if end <= start:
                raise ValueError("'end' must be greater than 'start'.")
            self.__end: float = end

        # current index of iterator
        self.__i: int = -1

    @property
    def start(self) -> float:
        """Sliding window start time in seconds."""
        return self.__start

    @property
    def end(self) -> float:
        """Sliding window end time in seconds."""
        return self.__end

    @property
    def step(self) -> float:
        """Sliding window step in seconds."""
        return self.__step

    @property
    def duration(self) -> float:
        """Sliding window duration in seconds."""
        return self.__duration

    def closest_frame(self, t: float) -> int:
        """Closest frame to timestamp.

        Parameters
        ----------
        t : float
            Timestamp, in seconds.

        Returns
        -------
        index : int
            Index of frame whose middle is the closest to `timestamp`

        """
        return int(np.rint(
            (t - self.__start - .5 * self.__duration) / self.__step
        ))

    '''
    def samples(self, from_duration: float, mode: Alignment = 'strict') -> int:
        """Number of frames

        Parameters
        ----------
        from_duration : float
            Duration in seconds.
        mode : {'strict', 'loose', 'center'}
            In 'strict' mode, computes the maximum number of consecutive frames
            that can be fitted into a segment with duration `from_duration`.
            In 'loose' mode, computes the maximum number of consecutive frames
            intersecting a segment with duration `from_duration`.
            In 'center' mode, computes the average number of consecutive frames
            where the first one is centered on the start time and the last one
            is centered on the end time of a segment with duration
            `from_duration`.

        """
        if mode == 'strict':
            return int(np.floor((from_duration - self.duration) / self.step)) + 1

        elif mode == 'loose':
            return int(np.floor((from_duration + self.duration) / self.step))

        elif mode == 'center':
            return int(np.rint((from_duration / self.step)))
    '''


    def segmentToRange(self, segment: Segment) -> Tuple[int, int]:
        warnings.warn("Deprecated in favor of `segment_to_range`",
                      DeprecationWarning)
        return self.segment_to_range(segment)

    def segment_to_range(self, segment: Segment) -> Tuple[int, int]:
        """Convert segment to 0-indexed frame range

        Parameters
        ----------
        segment : Segment

        Returns
        -------
        i0 : int
            Index of first frame
        n : int
            Number of frames

        Examples
        --------

            >>> window = SlidingWindow()
            >>> print window.segment_to_range(Segment(10, 15))
            i0, n

        """
        # find closest frame to segment start
        i0 = self.closest_frame(segment.start)

        # number of steps to cover segment duration
        n = int(segment.duration / self.step) + 1

        return i0, n

    def rangeToSegment(self, i0: int, n: int) -> Segment:
        warnings.warn("This is deprecated in favor of `range_to_segment`",
                      DeprecationWarning)
        return self.range_to_segment(i0, n)

    def range_to_segment(self, i0: int, n: int) -> Segment:
        """Convert 0-indexed frame range to segment

        Each frame represents a unique segment of duration 'step', centered on
        the middle of the frame.

        The very first frame (i0 = 0) is the exception. It is extended to the
        sliding window start time.

        Parameters
        ----------
        i0 : int
            Index of first frame
        n : int
            Number of frames

        Returns
        -------
        segment : Segment

        Examples
        --------

            >>> window = SlidingWindow()
            >>> print window.range_to_segment(3, 2)
            [ --> ]

        """

        # frame start time
        # start = self.start + i0 * self.step
        # frame middle time
        # start += .5 * self.duration
        # subframe start time
        # start -= .5 * self.step
        start = self.__start + (i0 - .5) * self.__step + .5 * self.__duration
        duration = n * self.__step
        end = start + duration

        # extend segment to the beginning of the timeline
        if i0 == 0:
            start = self.start

        return Segment(start, end)

    def samplesToDuration(self, nSamples: int) -> float:
        warnings.warn("This is deprecated in favor of `samples_to_duration`",
                      DeprecationWarning)
        return self.samples_to_duration(nSamples)

    def samples_to_duration(self, n_samples: int) -> float:
        """Returns duration of samples"""
        return self.range_to_segment(0, n_samples).duration

    def durationToSamples(self, duration: float) -> int:
        warnings.warn("This is deprecated in favor of `duration_to_samples`",
                      DeprecationWarning)
        return self.duration_to_samples(duration)

    def duration_to_samples(self, duration: float) -> int:
        """Returns samples in duration"""
        return self.segment_to_range(Segment(0, duration))[1]

    def __getitem__(self, i: int) -> Segment:
        """
        Parameters
        ----------
        i : int
            Index of sliding window position

        Returns
        -------
        segment : :class:`Segment`
            Sliding window at ith position

        """

        # window start time at ith position
        start = self.__start + i * self.__step

        # in case segment starts after the end,
        # return an empty segment
        if start >= self.__end:
            return None

        return Segment(start=start, end=start + self.__duration)

    def next(self) -> Segment:
        return self.__next__()

    def __next__(self) -> Segment:
        self.__i += 1
        window = self[self.__i]

        if window:
            return window
        else:
            raise StopIteration()

    def __iter__(self) -> 'SlidingWindow':
        """Sliding window iterator

        Use expression 'for segment in sliding_window'

        Examples
        --------

        >>> window = SlidingWindow(end=0.1)
        >>> for segment in window:
        ...     print(segment)
        [ 00:00:00.000 -->  00:00:00.030]
        [ 00:00:00.010 -->  00:00:00.040]
        [ 00:00:00.020 -->  00:00:00.050]
        [ 00:00:00.030 -->  00:00:00.060]
        [ 00:00:00.040 -->  00:00:00.070]
        [ 00:00:00.050 -->  00:00:00.080]
        [ 00:00:00.060 -->  00:00:00.090]
        [ 00:00:00.070 -->  00:00:00.100]
        [ 00:00:00.080 -->  00:00:00.110]
        [ 00:00:00.090 -->  00:00:00.120]
        """

        # reset iterator index
        self.__i = -1
        return self

    def __len__(self) -> int:
        """Number of positions

        Equivalent to len([segment for segment in window])

        Returns
        -------
        length : int
            Number of positions taken by the sliding window
            (from start times to end times)

        """
        if np.isinf(self.__end):
            raise ValueError('infinite sliding window.')

        # start looking for last position
        # based on frame closest to the end
        i = self.closest_frame(self.__end)

        while (self[i]):
            i += 1
        length = i

        return length

    def copy(self) -> 'SlidingWindow':
        """Duplicate sliding window"""
        duration = self.duration
        step = self.step
        start = self.start
        end = self.end
        sliding_window = self.__class__(
            duration=duration, step=step, start=start, end=end
        )
        return sliding_window

    def __call__(self,
                 support: Union[Segment, 'Timeline'],
                 align_last: bool = False) -> Iterable[Segment]:
        """Slide window over support

        Parameter
        ---------
        support : Segment or Timeline
            Support on which to slide the window.
        align_last : bool, optional
            Yield a final segment so that it aligns exactly with end of support.

        Yields
        ------
        chunk : Segment

        Example
        -------
        >>> window = SlidingWindow(duration=2., step=1.)
        >>> for chunk in window(Segment(3, 7.5)):
        ...     print(tuple(chunk))
        (3.0, 5.0)
        (4.0, 6.0)
        (5.0, 7.0)
        >>> for chunk in window(Segment(3, 7.5), align_last=True):
        ...     print(tuple(chunk))
        (3.0, 5.0)
        (4.0, 6.0)
        (5.0, 7.0)
        (5.5, 7.5)
        """

        from pyannote.core import Timeline
        if isinstance(support, Timeline):
            segments = support

        elif isinstance(support, Segment):
            segments = Timeline(segments=[support])

        else:
            msg = (
                f'"support" must be either a Segment or a Timeline '
                f'instance (is {type(support)})'
            )
            raise TypeError(msg)

        for segment in segments:

            if segment.duration < self.duration:
                continue

            window = SlidingWindow(duration=self.duration,
                                   step=self.step,
                                   start=segment.start,
                                   end=segment.end)

            for s in window:
                # ugly hack to account for floating point imprecision
                if s in segment:
                    yield s
                    last = s

            if align_last and last.end < segment.end:
                yield Segment(start=segment.end - self.duration,
                              end=segment.end)

class SlidingWindowFeature(np.lib.mixins.NDArrayOperatorsMixin):
    """Periodic feature vectors

    Parameters
    ----------
    data : (n_frames, n_features) numpy array
    sliding_window : SlidingWindow
    labels : list, optional
        Textual description of each dimension.
    """

    def __init__(
        self, data: np.ndarray, sliding_window: SlidingWindow, labels: List[Text] = None
    ):
        self.sliding_window: SlidingWindow = sliding_window
        self.data = data
        self.labels = labels
        self.__i: int = -1

    def __len__(self):
        """Number of feature vectors"""
        return self.data.shape[0]

    @property
    def extent(self):
        return self.sliding_window.range_to_segment(0, len(self))

    @property
    def dimension(self):
        """Dimension of feature vectors"""
        return self.data.shape[1]

    def getNumber(self):
        warnings.warn("This is deprecated in favor of `__len__`", DeprecationWarning)
        return self.data.shape[0]

    def getDimension(self):
        warnings.warn(
            "This is deprecated in favor of `dimension` property", DeprecationWarning
        )
        return self.dimension

    def getExtent(self):
        warnings.warn(
            "This is deprecated in favor of `extent` property", DeprecationWarning
        )
        return self.extent

    def __getitem__(self, i: int) -> np.ndarray:
        """Get ith feature vector"""
        return self.data[i]

    def __iter__(self):
        self.__i = -1
        return self

    def __next__(self) -> Tuple[Segment, np.ndarray]:
        self.__i += 1
        try:
            return self.sliding_window[self.__i], self.data[self.__i]
        except IndexError as e:
            raise StopIteration()

    def next(self):
        return self.__next__()

    def iterfeatures(
        self, window: Optional[bool] = False
    ) -> Iterator[Union[Tuple[np.ndarray, Segment], np.ndarray]]:
        """Feature vector iterator

        Parameters
        ----------
        window : bool, optional
            When True, yield both feature vector and corresponding window.
            Default is to only yield feature vector

        """
        n_samples = self.data.shape[0]
        for i in range(n_samples):
            if window:
                yield self.data[i], self.sliding_window[i]
            else:
                yield self.data[i]


    def _repr_png_(self):
        from .notebook import MATPLOTLIB_IS_AVAILABLE, MATPLOTLIB_WARNING

        if not MATPLOTLIB_IS_AVAILABLE:
            warnings.warn(MATPLOTLIB_WARNING.format(klass=self.__class__.__name__))
            return None

        from .notebook import repr_feature

        return repr_feature(self)

    _HANDLED_TYPES = (np.ndarray, numbers.Number)

    def __array__(self) -> np.ndarray:
        return self.data

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        out = kwargs.get("out", ())
        for x in inputs + out:
            # Only support operations with instances of _HANDLED_TYPES.
            # Use SlidingWindowFeature instead of type(self) for isinstance to
            # allow subclasses that don't override __array_ufunc__ to
            # handle SlidingWindowFeature objects.
            if not isinstance(x, self._HANDLED_TYPES + (SlidingWindowFeature,)):
                return NotImplemented

        # Defer to the implementation of the ufunc on unwrapped values.
        inputs = tuple(
            x.data if isinstance(x, SlidingWindowFeature) else x for x in inputs
        )
        if out:
            kwargs["out"] = tuple(
                x.data if isinstance(x, SlidingWindowFeature) else x for x in out
            )
        data = getattr(ufunc, method)(*inputs, **kwargs)

        if type(data) is tuple:
            # multiple return values
            return tuple(
                type(self)(x, self.sliding_window, labels=self.labels) for x in data
            )
        elif method == "at":
            # no return value
            return None
        else:
            # one return value
            return type(self)(data, self.sliding_window, labels=self.labels)




segment/utils.py
leohuang2013/pyannote-audio_speaker-diarization_cpp