File size: 1,117 Bytes
4ba35bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
"""
pyannote diarizer
"""
import matplotlib.pyplot as plt
from pyannote.audio import Pipeline
from pyannote.core import notebook
from diarizers.diarizer import Diarizer
class PyannoteDiarizer(Diarizer):
def __init__(self, audio_path: str):
"""
Pyannote diarizer class
Note: pyannote does not currently support defining the number of speakers, this
functionality might be supported later in the future
Args:
audio_path (str): the path to the audio file
Params:
diarization (Annotations): the output of the diarization algorithm
"""
self.audio_path = audio_path
self.diarization = None
def diarize_audio(self):
pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization")
self.diarization = pipeline({'audio': self.audio_path})
def get_diarization_figure(self) -> plt.gcf:
if not self.diarization:
self.diarize_audio()
figure, ax = plt.subplots()
notebook.plot_annotation(self.diarization, ax=ax, time=True, legend=True)
return plt.gcf()
|