first working version

app.py

@@ -0,0 +1,296 @@
+#!/usr/bin/env python3
+#
+# Copyright 2024 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# See LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# References:
+# https://gradio.app/docs/#dropdown
+
+import logging
+import os
+import time
+import uuid
+from datetime import datetime
+
+import gradio as gr
+import torch
+
+from model import (
+ embedding2models,
+ get_speaker_diarization,
+ read_wave,
+ speaker_segmentation_models,
+)
+
+embedding_frameworks = list(embedding2models.keys())
+
+
+def MyPrint(s):
+ now = datetime.now()
+ date_time = now.strftime("%Y-%m-%d %H:%M:%S.%f")
+ print(f"{date_time}: {s}")
+
+
+def convert_to_wav(in_filename: str) -> str:
+ """Convert the input audio file to a wave file"""
+ out_filename = str(uuid.uuid4())
+ out_filename = f"{in_filename}.wav"
+
+ MyPrint(f"Converting '{in_filename}' to '{out_filename}'")
+ _ = os.system(
+ f"ffmpeg -hide_banner -loglevel error -i '{in_filename}' -ar 16000 -ac 1 '{out_filename}' -y"
+ )
+
+ return out_filename
+
+
+def build_html_output(s: str, style: str = "result_item_success"):
+ return f"""
+ <div class='result'>
+ <div class='result_item {style}'>
+ {s}
+ </div>
+ </div>
+ """
+
+
+def process_uploaded_file(
+ embedding_framework: str,
+ embedding_model: str,
+ speaker_segmentation_model: str,
+ input_num_speakers: str,
+ input_threshold: str,
+ in_filename: str,
+):
+ if in_filename is None or in_filename == "":
+ return "", build_html_output(
+ "Please first upload a file and then click "
+ 'the button "submit for recognition"',
+ "result_item_error",
+ )
+
+ try:
+ input_num_speakers = int(input_num_speakers)
+ except ValueError:
+ return "", build_html_output(
+ "Please set a valid number of speakers",
+ "result_item_error",
+ )
+
+ try:
+ input_threshold = float(input_threshold)
+ if input_threshold < 0 or input_threshold < 10:
+ raise ValueError("")
+ except ValueError:
+ return "", build_html_output(
+ "Please set a valid threshold between (0, 10)",
+ "result_item_error",
+ )
+
+ MyPrint(f"Processing uploaded file: {in_filename}")
+ try:
+ return process(
+ in_filename=in_filename,
+ embedding_framework=embedding_framework,
+ embedding_model=embedding_model,
+ speaker_segmentation_model=speaker_segmentation_model,
+ input_num_speakers=input_num_speakers,
+ input_threshold=input_threshold,
+ )
+ except Exception as e:
+ MyPrint(str(e))
+ return "", build_html_output(str(e), "result_item_error")
+
+
+@torch.no_grad()
+def process(
+ embedding_framework: str,
+ embedding_model: str,
+ speaker_segmentation_model: str,
+ input_num_speakers: str,
+ input_threshold: str,
+ in_filename: str,
+):
+ MyPrint(f"embedding_framework: {embedding_framework}")
+ MyPrint(f"embedding_model: {embedding_model}")
+ MyPrint(f"speaker_segmentation_model: {speaker_segmentation_model}")
+ MyPrint(f"input_num_speakers: {input_num_speakers}")
+ MyPrint(f"input_threshold: {input_threshold}")
+ MyPrint(f"in_filename: {in_filename}")
+
+ filename = convert_to_wav(in_filename)
+
+ now = datetime.now()
+ date_time = now.strftime("%Y-%m-%d %H:%M:%S.%f")
+ MyPrint(f"Started at {date_time}")
+
+ start = time.time()
+
+ sd = get_speaker_diarization(
+ segmentation=speaker_segmentation_model,
+ embedding_model=embedding_model,
+ num_clusters=input_num_speakers,
+ threshold=input_threshold,
+ )
+
+ audio = read_wave(filename)[0]
+ segments = sd.process(audio).sort_by_start_time()
+ s = ""
+ for seg in segments:
+ s += f"{seg.start:.3f} -- {seg.end:.3f} speaker_{seg.speaker:02d}\n"
+
+ date_time = now.strftime("%Y-%m-%d %H:%M:%S.%f")
+ end = time.time()
+
+ duration = audio.shape[0] / sd.sample_rate
+ rtf = (end - start) / duration
+
+ MyPrint(f"Finished at {date_time} s. Elapsed: {end - start: .3f} s")
+
+ info = f"""
+ Wave duration : {duration: .3f} s <br/>
+ Processing time: {end - start: .3f} s <br/>
+ RTF: {end - start: .3f}/{duration: .3f} = {rtf:.3f} <br/>
+ """
+ if rtf > 1:
+ info += (
+ "<br/>We are loading the model for the first run. "
+ "Please run again to measure the real RTF.<br/>"
+ )
+
+ MyPrint(info)
+ MyPrint(f"\nembedding_model: {embedding_model}\nSegments: {s}")
+
+ return s, build_html_output(info)
+
+
+title = "# Speaker diarization with Next-gen Kaldi"
+description = """
+This space shows how to do speaker diarization with Next-gen Kaldi.
+
+It is running on CPU within a docker container provided by Hugging Face.
+
+See more information by visiting
+<https://k2-fsa.github.io/sherpa/onnx/speaker-diarization/index.html>
+
+If you want to try it on Android, please download pre-built Android
+APKs for speaker diarzation by visiting
+<https://k2-fsa.github.io/sherpa/onnx/speaker-diarization/android.html>
+"""
+
+# css style is copied from
+# https://huggingface.co/spaces/alphacep/asr/blob/main/app.py#L113
+css = """
+.result {display:flex;flex-direction:column}
+.result_item {padding:15px;margin-bottom:8px;border-radius:15px;width:100%}
+.result_item_success {background-color:mediumaquamarine;color:white;align-self:start}
+.result_item_error {background-color:#ff7070;color:white;align-self:start}
+"""
+
+
+def update_embedding_model_dropdown(framework: str):
+ if framework in embedding2models:
+ choices = embedding2models[framework]
+ return gr.Dropdown(
+ choices=choices,
+ value=choices[0],
+ interactive=True,
+ )
+
+ raise ValueError(f"Unsupported framework: {framework}")
+
+
+demo = gr.Blocks(css=css)
+
+
+with demo:
+ gr.Markdown(title)
+
+ embedding_framework_choices = list(embedding2models.keys())
+ embedding_framework_radio = gr.Radio(
+ label="Speaker embedding frameworks",
+ choices=embedding_framework_choices,
+ value=embedding_framework_choices[0],
+ )
+
+ embedding_model_dropdown = gr.Dropdown(
+ choices=embedding2models[embedding_framework_choices[0]],
+ label="Select a speaker embedding model",
+ value=embedding2models[embedding_framework_choices[0]][0],
+ )
+
+ embedding_framework_choices.change(
+ update_embedding_model_dropdown,
+ inputs=embedding_framework_radio,
+ outputs=embedding_model_dropdown,
+ )
+
+ speaker_segmentation_model_dropdown = gr.Dropdown(
+ choices=speaker_segmentation_models,
+ label="Select a speaker segmentation model",
+ value=speaker_segmentation_models[0],
+ )
+
+ input_num_speakers = gr.Textbox(
+ label="Number of speakers",
+ info="Number of speakers",
+ lines=1,
+ max_lines=1,
+ value="0",
+ placeholder="Specify number of speakers in the test file",
+ )
+
+ input_threshold = gr.Textbox(
+ label="Clustering threshold",
+ info="Threshold for clustering",
+ lines=1,
+ max_lines=1,
+ value="0.5",
+ placeholder="Clustering for threshold",
+ )
+
+ with gr.Tabs():
+ with gr.TabItem("Upload from disk"):
+ uploaded_file = gr.Audio(
+ sources=["upload"], # Choose between "microphone", "upload"
+ type="filepath",
+ label="Upload from disk",
+ )
+ upload_button = gr.Button("Submit for speaker diarization")
+ uploaded_output = gr.Textbox(label="Result from uploaded file")
+ uploaded_html_info = gr.HTML(label="Info")
+
+ upload_button.click(
+ process_uploaded_file,
+ inputs=[
+ embedding_framework_radio,
+ embedding_model_dropdown,
+ speaker_segmentation_model_dropdown,
+ input_num_speakers,
+ input_threshold,
+ uploaded_file,
+ ],
+ outputs=[uploaded_output, uploaded_html_info],
+ )
+
+ gr.Markdown(description)
+
+if __name__ == "__main__":
+ formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
+
+ logging.basicConfig(format=formatter, level=logging.INFO)
+
+ demo.launch()

model.py

@@ -0,0 +1,152 @@
+# Copyright 2024 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# See LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import wave
+from typing import List, Tuple
+
+import numpy as np
+import sherpa_onnx
+from huggingface_hub import hf_hub_download
+
+
+def read_wave(wave_filename: str) -> Tuple[np.ndarray, int]:
+ """
+ Args:
+ wave_filename:
+ Path to a wave file. It should be single channel and each sample should
+ be 16-bit. Its sample rate does not need to be 16kHz.
+ Returns:
+ Return a tuple containing:
+ - A 1-D array of dtype np.float32 containing the samples, which are
+ normalized to the range [-1, 1].
+ - sample rate of the wave file
+ """
+
+ with wave.open(wave_filename) as f:
+ assert f.getnchannels() == 1, f.getnchannels()
+ assert f.getsampwidth() == 2, f.getsampwidth() # it is in bytes
+ num_samples = f.getnframes()
+ samples = f.readframes(num_samples)
+ samples_int16 = np.frombuffer(samples, dtype=np.int16)
+ samples_float32 = samples_int16.astype(np.float32)
+
+ samples_float32 = samples_float32 / 32768
+ return samples_float32, f.getframerate()
+
+
+def _get_nn_model_filename(
+ repo_id: str,
+ filename: str,
+ subfolder: str = ".",
+) -> str:
+ nn_model_filename = hf_hub_download(
+ repo_id=repo_id,
+ filename=filename,
+ subfolder=subfolder,
+ )
+ return nn_model_filename
+
+
+def get_speaker_segmentation_model(repo_id) -> List[str]:
+ assert repo_id in ("pyannote/segmentation-3.0",)
+
+ if repo_id == "pyannote/segmentation-3.0":
+ return _get_nn_model_filename(
+ repo_id="csukuangfj/sherpa-onnx-pyannote-segmentation-3-0",
+ filename="model.onnx",
+ )
+
+
+def get_speaker_embedding_model(model_name) -> List[str]:
+ assert (
+ model_name
+ in three_d_speaker_embedding_models
+ + nemo_speaker_embedding_models
+ + wespeaker_embedding_models
+ )
+
+ return _get_nn_model_filename(
+ repo_id="csukuangfj/speaker-embedding-models",
+ filename=model_name,
+ )
+
+
+def get_speaker_diarization(
+ segmentation_model: str, embedding_model: str, num_clusters: int, threshold: float
+):
+ segmentation = get_speaker_segmentation_model(segmentation_model)
+ embedding = get_speaker_embedding_model(embedding_model)
+
+ config = sherpa_onnx.OfflineSpeakerDiarizationConfig(
+ segmentation=sherpa_onnx.OfflineSpeakerSegmentationModelConfig(
+ pyannote=sherpa_onnx.OfflineSpeakerSegmentationPyannoteModelConfig(
+ model=segmentation
+ ),
+ ),
+ embedding=sherpa_onnx.SpeakerEmbeddingExtractorConfig(model=embedding),
+ clustering=sherpa_onnx.FastClusteringConfig(
+ num_clusters=num_clusters,
+ threshold=threshold,
+ ),
+ min_duration_on=0.3,
+ min_duration_off=0.5,
+ )
+ if not config.validate():
+ raise RuntimeError(
+ "Please check your config and make sure all required files exist"
+ )
+
+ return sherpa_onnx.OfflineSpeakerDiarization(config)
+ pass
+
+
+speaker_segmentation_models = ["pyannote/segmentation-3.0"]
+
+
+nemo_speaker_embedding_models = [
+ "nemo_en_speakerverification_speakernet.onnx",
+ "nemo_en_titanet_large.onnx",
+ "nemo_en_titanet_small.onnx",
+]
+
+three_d_speaker_embedding_models = [
+ "3dspeaker_speech_campplus_sv_en_voxceleb_16k.onnx",
+ "3dspeaker_speech_campplus_sv_zh-cn_16k-common.onnx",
+ "3dspeaker_speech_campplus_sv_zh_en_16k-common_advanced.onnx",
+ "3dspeaker_speech_eres2net_base_200k_sv_zh-cn_16k-common.onnx",
+ "3dspeaker_speech_eres2net_base_sv_zh-cn_3dspeaker_16k.onnx",
+ "3dspeaker_speech_eres2net_large_sv_zh-cn_3dspeaker_16k.onnx",
+ "3dspeaker_speech_eres2net_sv_en_voxceleb_16k.onnx",
+ "3dspeaker_speech_eres2net_sv_zh-cn_16k-common.onnx",
+ "3dspeaker_speech_eres2netv2_sv_zh-cn_16k-common.onnx",
+]
+wespeaker_embedding_models = [
+ "wespeaker_en_voxceleb_CAM++.onnx",
+ "wespeaker_en_voxceleb_CAM++_LM.onnx",
+ "wespeaker_en_voxceleb_resnet152_LM.onnx",
+ "wespeaker_en_voxceleb_resnet221_LM.onnx",
+ "wespeaker_en_voxceleb_resnet293_LM.onnx",
+ "wespeaker_en_voxceleb_resnet34.onnx",
+ "wespeaker_en_voxceleb_resnet34_LM.onnx",
+ "wespeaker_zh_cnceleb_resnet34.onnx",
+ "wespeaker_zh_cnceleb_resnet34_LM.onnx",
+]
+
+embedding2models = {
+ "3D-Speaker": three_d_speaker_embedding_models,
+ "NeMo": nemo_speaker_embedding_models,
+ "WeSpeaker": wespeaker_embedding_models,
+}

requirements.txt

@@ -0,0 +1,5 @@
+huggingface_hub
+
+#https://huggingface.co/csukuangfj/sherpa-onnx-wheels/resolve/main/sherpa_onnx-1.9.26-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+
+sherpa-onnx>=1.10.28