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