hicustomer/pyannote-speaker-diarization

1from typing import Dict2from pyannote.audio import Pipeline3from io import BytesIO4import torch5import torchaudio6 7 8class EndpointHandler:9    def __init__(self, path=""):10        # load the model11        self.pipeline = Pipeline.from_pretrained("config.yaml")12 13    def __call__(self, data: Dict[str, bytes]) -> Dict[str, str]:14        """15        Args:16            data (:obj:):17                includes the deserialized audio file as bytes18        Return:19            A :obj:`dict`:. base64 encoded image20        """21        # process input22        inputs = data.pop("inputs", data)23        parameters = data.pop("parameters", None)  #  min_speakers=2, max_speakers=524 25        waveform, sample_rate = torchaudio.load(BytesIO(inputs))26        pyannote_input = {"waveform": waveform, "sample_rate": sample_rate}27 28        # apply pretrained pipeline29        # pass inputs with all kwargs in data30        if parameters is not None:31            diarization = self.pipeline(pyannote_input, **parameters)32        else:33            diarization = self.pipeline(pyannote_input)34 35        # postprocess the prediction36        processed_diarization = [37            {"label": str(label), "start": str(segment.start), "stop": str(segment.end)}38            for segment, _, label in diarization.itertracks(yield_label=True)39        ]40 41        return {"diarization": processed_diarization}42