KIFF
/

pyannote-speaker-diarization-endpoint

Voice Activity Detection

pyannote-audio-pipeline

speaker-diarization

speaker-change-detection

overlapped-speech-detection

Inference Endpoints

Model card Files Files and versions Community

pyannote-speaker-diarization-endpoint / handler.py

KIFF's picture

Update handler.py

7c12535 11 months ago

raw history blame contribute delete

1.64 kB

	from typing import Dict
	from pyannote.audio import Pipeline
	import torch
	import base64
	import numpy as np

	SAMPLE_RATE = 16000

	class EndpointHandler():
	def __init__(self, path=""):
	# load the model
	self.pipeline = Pipeline.from_pretrained("KIFF/pyannote-speaker-diarization-endpoint")

	def __call__(self, data: Dict[str, bytes]) -> Dict[str, str]:
	"""
	Args:
	data (:obj:):
	includes the deserialized audio file as bytes
	Return:
	A :obj:`dict`:. base64 encoded image
	"""
	# process input
	inputs = data.pop("inputs", data)
	parameters = data.pop("parameters", None) # min_speakers=2, max_speakers=5

	# decode the base64 audio data
	audio_data = base64.b64decode(inputs)
	audio_nparray = np.frombuffer(audio_data, dtype=np.int16)

	# prepare pynannote input
	audio_tensor= torch.from_numpy(audio_nparray).float().unsqueeze(0)
	pyannote_input = {"waveform": audio_tensor, "sample_rate": SAMPLE_RATE}

	# apply pretrained pipeline
	# pass inputs with all kwargs in data
	if parameters is not None:
	diarization = self.pipeline(pyannote_input, **parameters)
	else:
	diarization = self.pipeline(pyannote_input)

	# postprocess the prediction
	processed_diarization = [
	{"label": str(label), "start": str(segment.start), "stop": str(segment.end)}
	for segment, _, label in diarization.itertracks(yield_label=True)
	]

	return {"diarization": processed_diarization}