DroolingPanda/speaker-diarization-3.1

1# MIT License2#3# Copyright (c) 2023 CNRS4#5# Permission is hereby granted, free of charge, to any person obtaining a copy6# of this software and associated documentation files (the "Software"), to deal7# in the Software without restriction, including without limitation the rights8# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell9# copies of the Software, and to permit persons to whom the Software is10# furnished to do so, subject to the following conditions:11#12# The above copyright notice and this permission notice shall be included in all13# copies or substantial portions of the Software.14#15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR16# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,17# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE18# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER19# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE21# SOFTWARE.22 23 24from pyannote.audio import Pipeline, Audio25import torch26 27 28class EndpointHandler:29    def __init__(self, path=""):30        # initialize pretrained pipeline31        self._pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")32 33        # send pipeline to GPU if available34        if torch.cuda.is_available():35            self._pipeline.to(torch.device("cuda"))36 37        # initialize audio reader38        self._io = Audio()39 40    def __call__(self, data):41        inputs = data.pop("inputs", data)42        waveform, sample_rate = self._io(inputs)43 44        parameters = data.pop("parameters", dict())45        diarization = self.pipeline(46            {"waveform": waveform, "sample_rate": sample_rate}, **parameters47        )48 49        processed_diarization = [50            {51                "speaker": speaker,52                "start": f"{turn.start:.3f}",53                "end": f"{turn.end:.3f}",54            }55            for turn, _, speaker in diarization.itertracks(yield_label=True)56        ]57 58        return {"diarization": processed_diarization}59