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

scripts/apk/build-apk-speaker-diarization.sh.in

@@ -31,15 +31,24 @@ log "====================x86===================="
 
 mkdir -p apks
 
+{% for model in model_list %}
+
 pushd ./android/SherpaOnnxSpeakerDiarization/app/src/main/assets/
 
-curl -SL -O https://github.com/k2-fsa/sherpa-onnx/releases/download/speaker-segmentation-models/sherpa-onnx-pyannote-segmentation-3-0.tar.bz2
-tar xvf sherpa-onnx-pyannote-segmentation-3-0.tar.bz2
-rm sherpa-onnx-pyannote-segmentation-3-0.tar.bz2
-mv sherpa-onnx-pyannote-segmentation-3-0/model.onnx segmentation.onnx
-rm -rf sherpa-onnx-pyannote-segmentation-3-0
+ls -lh
 
-curl -SL -O https://github.com/k2-fsa/sherpa-onnx/releases/download/speaker-recongition-models/3dspeaker_speech_eres2net_base_sv_zh-cn_3dspeaker_16k.onnx
+model_name={{ model.model_name }}
+short_name={{ model.short_name }}
+
+curl -SL -O https://github.com/k2-fsa/sherpa-onnx/releases/download/speaker-segmentation-models/$model_name.tar.bz2
+tar xvf $model_name.tar.bz2
+rm $model_name.tar.bz2
+mv $model_name/model.onnx segmentation.onnx
+rm -rf $model_name
+
+if [ ! -f 3dspeaker_speech_eres2net_base_sv_zh-cn_3dspeaker_16k.onnx ]; then
+  curl -SL -O https://github.com/k2-fsa/sherpa-onnx/releases/download/speaker-recongition-models/3dspeaker_speech_eres2net_base_sv_zh-cn_3dspeaker_16k.onnx
+fi
 
 echo "pwd: $PWD"
 ls -lh
@@ -65,9 +74,13 @@ for arch in arm64-v8a armeabi-v7a x86_64 x86; do
   ./gradlew build
   popd
 
-  mv android/SherpaOnnxSpeakerDiarization/app/build/outputs/apk/debug/app-debug.apk ./apks/sherpa-onnx-${SHERPA_ONNX_VERSION}-$arch-speaker-diarization-pyannote_audio-3dspeaker.apk
+  mv android/SherpaOnnxSpeakerDiarization/app/build/outputs/apk/debug/app-debug.apk ./apks/sherpa-onnx-${SHERPA_ONNX_VERSION}-$arch-speaker-diarization-$short_name-3dspeaker.apk
   ls -lh apks
   rm -v ./android/SherpaOnnxSpeakerDiarization/app/src/main/jniLibs/$arch/*.so
 done
 
+rm -rf ./android/SherpaOnnxSpeakerDiarization/app/src/main/assets/segmentation.onnx
+
+{% endfor %}
+
 ls -lh apks

scripts/apk/generate-speaker-diarization-apk-script.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+
+import argparse
+from dataclasses import dataclass
+from typing import List
+
+import jinja2
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--total",
+        type=int,
+        default=1,
+        help="Number of runners",
+    )
+    parser.add_argument(
+        "--index",
+        type=int,
+        default=0,
+        help="Index of the current runner",
+    )
+    return parser.parse_args()
+
+
+@dataclass
+class SpeakerSegmentationModel:
+    model_name: str
+    short_name: str = ""
+
+
+def get_models() -> List[SpeakerSegmentationModel]:
+    models = [
+        SpeakerSegmentationModel(
+            model_name="sherpa-onnx-pyannote-segmentation-3-0",
+            short_name="pyannote_audio",
+        ),
+        SpeakerSegmentationModel(
+            model_name="sherpa-onnx-reverb-diarization-v1",
+            short_name="revai_v1",
+        ),
+    ]
+
+    return models
+
+
+def main():
+    args = get_args()
+    index = args.index
+    total = args.total
+    assert 0 <= index < total, (index, total)
+
+    all_model_list = get_models()
+
+    num_models = len(all_model_list)
+
+    num_per_runner = num_models // total
+    if num_per_runner <= 0:
+        raise ValueError(f"num_models: {num_models}, num_runners: {total}")
+
+    start = index * num_per_runner
+    end = start + num_per_runner
+
+    remaining = num_models - args.total * num_per_runner
+
+    print(f"{index}/{total}: {start}-{end}/{num_models}")
+
+    d = dict()
+    d["model_list"] = all_model_list[start:end]
+    if index < remaining:
+        s = args.total * num_per_runner + index
+        d["model_list"].append(all_model_list[s])
+        print(f"{s}/{num_models}")
+
+    filename_list = ["./build-apk-speaker-diarization.sh"]
+    for filename in filename_list:
+        environment = jinja2.Environment()
+        with open(f"{filename}.in") as f:
+            s = f.read()
+        template = environment.from_string(s)
+
+        s = template.render(**d)
+        with open(filename, "w") as f:
+            print(s, file=f)
+
+
+if __name__ == "__main__":
+    main()