Support https://huggingface.co/Revai/reverb-diarization-v1 (#1437)

scripts/pyannote/segmentation/export-onnx.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
+# Copyright      2024  Xiaomi Corp.        (authors: Fangjun Kuang)
 
+import os
 from typing import Any, Dict
 
 import onnx
@@ -35,6 +37,8 @@ def add_meta_data(filename: str, meta_data: Dict[str, Any]):
 def main():
     # You can download ./pytorch_model.bin from
     # https://hf-mirror.com/csukuangfj/pyannote-models/tree/main/segmentation-3.0
+    # or from
+    # https://huggingface.co/Revai/reverb-diarization-v1/tree/main
     pt_filename = "./pytorch_model.bin"
     model = Model.from_pretrained(pt_filename)
     model.eval()
@@ -94,6 +98,22 @@ def main():
     receptive_field_size = int(model.receptive_field.duration * 16000)
     receptive_field_shift = int(model.receptive_field.step * 16000)
 
+    is_revai = os.getenv("SHERPA_ONNX_IS_REVAI", "")
+    if is_revai == "":
+        url_1 = "https://huggingface.co/pyannote/segmentation-3.0"
+        url_2 = "https://huggingface.co/csukuangfj/pyannote-models/tree/main/segmentation-3.0"
+        license_url = (
+            "https://huggingface.co/pyannote/segmentation-3.0/blob/main/LICENSE"
+        )
+        model_author = "pyannote-audio"
+    else:
+        url_1 = "https://huggingface.co/Revai/reverb-diarization-v1"
+        url_2 = "https://huggingface.co/csukuangfj/sherpa-onnx-reverb-diarization-v1"
+        license_url = (
+            "https://huggingface.co/Revai/reverb-diarization-v1/blob/main/LICENSE"
+        )
+        model_author = "Revai"
+
     meta_data = {
         "num_speakers": len(model.specifications.classes),
         "powerset_max_classes": model.specifications.powerset_max_classes,
@@ -104,11 +124,11 @@ def main():
         "receptive_field_shift": receptive_field_shift,
         "model_type": "pyannote-segmentation-3.0",
         "version": "1",
-        "model_author": "pyannote",
+        "model_author": model_author,
         "maintainer": "k2-fsa",
-        "url_1": "https://huggingface.co/pyannote/segmentation-3.0",
-        "url_2": "https://huggingface.co/csukuangfj/pyannote-models/tree/main/segmentation-3.0",
-        "license": "https://huggingface.co/pyannote/segmentation-3.0/blob/main/LICENSE",
+        "url_1": url_1,
+        "url_2": url_2,
+        "license": license_url,
     }
     add_meta_data(filename=filename, meta_data=meta_data)
 

scripts/pyannote/segmentation/preprocess.sh

@@ -1,4 +1,5 @@
 #!/usr/bin/env bash
+# Copyright      2024  Xiaomi Corp.        (authors: Fangjun Kuang)
 
 
 python3 -m onnxruntime.quantization.preprocess --input model.onnx --output tmp.preprocessed.onnx

scripts/pyannote/segmentation/run-revai.sh

@@ -0,0 +1,48 @@
+#!/usr/bin/env bash
+# Copyright    2024  Xiaomi Corp.        (authors: Fangjun Kuang)
+
+export SHERPA_ONNX_IS_REVAI=1
+
+set -ex
+function install_pyannote() {
+  pip install pyannote.audio onnx onnxruntime
+}
+
+function download_test_files() {
+  curl -SL -O https://huggingface.co/Revai/reverb-diarization-v1/resolve/main/pytorch_model.bin
+  curl -SL -O https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/lei-jun-test.wav
+}
+
+install_pyannote
+download_test_files
+
+./export-onnx.py
+./preprocess.sh
+
+echo "----------torch----------"
+./vad-torch.py
+
+echo "----------onnx model.onnx----------"
+./vad-onnx.py --model ./model.onnx --wav ./lei-jun-test.wav
+
+echo "----------onnx model.int8.onnx----------"
+./vad-onnx.py --model ./model.int8.onnx --wav ./lei-jun-test.wav
+
+curl -SL -O https://huggingface.co/Revai/reverb-diarization-v1/resolve/main/LICENSE
+
+cat >README.md << EOF
+# Introduction
+
+Models in this file are converted from
+https://huggingface.co/Revai/reverb-diarization-v1/tree/main
+
+Note that it is accessible under a non-commercial license.
+
+Please see ./LICENSE for details.
+
+See also
+https://www.rev.com/blog/speech-to-text-technology/introducing-reverb-open-source-asr-diarization
+
+EOF
+
+

scripts/pyannote/segmentation/speaker-diarization-torch.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+# Copyright      2024  Xiaomi Corp.        (authors: Fangjun Kuang)
 
 """
 Please refer to

scripts/pyannote/segmentation/vad-onnx.py

@@ -216,6 +216,8 @@ def main():
 
     is_active = classification[0] > onset
     start = None
+    if is_active:
+        start = 0
 
     scale = m.receptive_field_shift / m.sample_rate
     scale_offset = m.receptive_field_size / m.sample_rate * 0.5