chrissewell/whisperX
1
0
Fork 0
mirror of https://github.com/m-bain/whisperX.git synced 2025-07-01 18:17:27 -04:00
Code Issues Packages Projects Releases Wiki Activity

skeleton v2

Browse Source
This commit is contained in:
Max Bain
2023-03-30 05:31:57 +01:00
parent 1e7c2c337b
commit 18b63d46e2
53 changed files with 752 additions and 106949 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
whisperx/diarize.py
View File

@ -1,8 +1,23 @@
import numpy as np
import pandas as pd
from pyannote.audio import Pipeline
class DiarizationPipeline:
def __init__(
self,
model_name="pyannote/speaker-diarization@2.1",
use_auth_token=None,
):
self.model = Pipeline.from_pretrained(model_name, use_auth_token=use_auth_token)
def __call__(self, audio, min_speakers=None, max_speakers=None):
segments = self.model(audio, min_speakers=min_speakers, max_speakers=max_speakers)
diarize_df = pd.DataFrame(segments.itertracks(yield_label=True))
diarize_df['start'] = diarize_df[0].apply(lambda x: x.start)
diarize_df['end'] = diarize_df[0].apply(lambda x: x.end)
return diarize_df
def assign_word_speakers(diarize_df, result_segments, fill_nearest=False):
for seg in result_segments:
wdf = seg['word-segments']
if len(wdf['start'].dropna()) == 0: