Add Rust API for speaker diarization (#3370)

.github/scripts/test-rust.sh

@@ -18,6 +18,10 @@ rm -rf sr-data
 ./run-speaker-embedding-cosine-similarity.sh
 rm -f wespeaker_zh_cnceleb_resnet34.onnx fangjun-sr-1.wav fangjun-sr-2.wav leijun-sr-1.wav
 
+./run-offline-speaker-diarization.sh
+rm -rf sherpa-onnx-pyannote-segmentation-3-0
+rm -f 3dspeaker_speech_eres2net_base_sv_zh-cn_3dspeaker_16k.onnx 0-four-speakers-zh.wav
+
 ./run-vits-en.sh
 rm -rf vits-piper-en_US-amy-low
 

rust-api-examples/README.md

@@ -65,6 +65,7 @@ export RUSTFLAGS="-C link-arg=-Wl,-rpath,$SHERPA_ONNX_LIB_DIR"
 | 30 | [speaker_embedding_extractor](#example-30-speaker-embedding-extractor) | Compute a speaker embedding from a wave file |
 | 31 | [speaker_embedding_manager](#example-31-speaker-embedding-manager) | Register, search, verify, and remove speakers using embeddings |
 | 32 | [speaker_embedding_cosine_similarity](#example-32-speaker-embedding-cosine-similarity) | Compute cosine similarity from three speaker embeddings |
+| 33 | [offline_speaker_diarization](#example-33-offline-speaker-diarization) | Offline speaker diarization with pyannote segmentation and 3D-Speaker embeddings |
 
 ## Run it
 
@@ -270,3 +271,10 @@ to check the RPATH.
 ```bash
 ./run-speaker-embedding-cosine-similarity.sh
 ```
+
+
+### Example 33: Offline speaker diarization
+
+```bash
+./run-offline-speaker-diarization.sh
+```

rust-api-examples/examples/offline_speaker_diarization.rs

@@ -0,0 +1,46 @@
+use sherpa_onnx::{
+    FastClusteringConfig, OfflineSpeakerDiarization, OfflineSpeakerDiarizationConfig,
+    OfflineSpeakerSegmentationModelConfig, OfflineSpeakerSegmentationPyannoteModelConfig,
+    SpeakerEmbeddingExtractorConfig, Wave,
+};
+
+fn main() {
+    let config = OfflineSpeakerDiarizationConfig {
+        segmentation: OfflineSpeakerSegmentationModelConfig {
+            pyannote: OfflineSpeakerSegmentationPyannoteModelConfig {
+                model: Some("./sherpa-onnx-pyannote-segmentation-3-0/model.onnx".into()),
+            },
+            ..Default::default()
+        },
+        embedding: SpeakerEmbeddingExtractorConfig {
+            model: Some("./3dspeaker_speech_eres2net_base_sv_zh-cn_3dspeaker_16k.onnx".into()),
+            ..Default::default()
+        },
+        clustering: FastClusteringConfig {
+            num_clusters: 4,
+            ..Default::default()
+        },
+        ..Default::default()
+    };
+
+    let sd = OfflineSpeakerDiarization::create(&config)
+        .expect("Failed to initialize offline speaker diarization");
+
+    let wave = Wave::read("./0-four-speakers-zh.wav").expect("Failed to read wave");
+
+    assert_eq!(
+        sd.sample_rate(),
+        wave.sample_rate(),
+        "Unexpected sample rate"
+    );
+
+    let result = sd
+        .process(wave.samples())
+        .expect("Failed to do speaker diarization");
+    println!("Number of speakers: {}", result.num_speakers());
+    println!("Number of segments: {}", result.num_segments());
+
+    for s in result.sort_by_start_time() {
+        println!("{:.3} -- {:.3} speaker_{:02}", s.start, s.end, s.speaker);
+    }
+}

rust-api-examples/run-offline-speaker-diarization.sh

@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+set -ex
+
+if [ ! -f ./sherpa-onnx-pyannote-segmentation-3-0/model.onnx ]; then
+  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
+fi
+
+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
+
+if [ ! -f ./0-four-speakers-zh.wav ]; then
+  curl -SL -O https://github.com/k2-fsa/sherpa-onnx/releases/download/speaker-segmentation-models/0-four-speakers-zh.wav
+fi
+
+cargo run --example offline_speaker_diarization

sherpa-onnx/rust/sherpa-onnx-sys/src/lib.rs

@@ -13,6 +13,7 @@ extern "C" {
 
 pub mod audio_tagging;
 pub mod offline_asr;
+pub mod offline_speaker_diarization;
 pub mod online_asr;
 pub mod online_punctuation;
 pub mod speaker_embedding;
@@ -23,6 +24,7 @@ pub mod wave;
 
 pub use audio_tagging::*;
 pub use offline_asr::*;
+pub use offline_speaker_diarization::*;
 pub use online_asr::*;
 pub use online_punctuation::*;
 pub use speaker_embedding::*;

sherpa-onnx/rust/sherpa-onnx-sys/src/offline_speaker_diarization.rs

@@ -0,0 +1,98 @@
+#![allow(non_camel_case_types)]
+#![allow(non_snake_case)]
+#![allow(non_upper_case_globals)]
+
+use std::os::raw::{c_char, c_float};
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub struct OfflineSpeakerSegmentationPyannoteModelConfig {
+    pub model: *const c_char,
+}
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub struct OfflineSpeakerSegmentationModelConfig {
+    pub pyannote: OfflineSpeakerSegmentationPyannoteModelConfig,
+    pub num_threads: i32,
+    pub debug: i32,
+    pub provider: *const c_char,
+}
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub struct FastClusteringConfig {
+    pub num_clusters: i32,
+    pub threshold: c_float,
+}
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub struct OfflineSpeakerDiarizationConfig {
+    pub segmentation: OfflineSpeakerSegmentationModelConfig,
+    pub embedding: crate::speaker_embedding::SpeakerEmbeddingExtractorConfig,
+    pub clustering: FastClusteringConfig,
+    pub min_duration_on: c_float,
+    pub min_duration_off: c_float,
+}
+
+#[repr(C)]
+pub struct OfflineSpeakerDiarization {
+    _private: [u8; 0],
+}
+
+#[repr(C)]
+pub struct OfflineSpeakerDiarizationResult {
+    _private: [u8; 0],
+}
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub struct OfflineSpeakerDiarizationSegment {
+    pub start: c_float,
+    pub end: c_float,
+    pub speaker: i32,
+}
+
+extern "C" {
+    pub fn SherpaOnnxCreateOfflineSpeakerDiarization(
+        config: *const OfflineSpeakerDiarizationConfig,
+    ) -> *const OfflineSpeakerDiarization;
+
+    pub fn SherpaOnnxDestroyOfflineSpeakerDiarization(sd: *const OfflineSpeakerDiarization);
+
+    pub fn SherpaOnnxOfflineSpeakerDiarizationGetSampleRate(
+        sd: *const OfflineSpeakerDiarization,
+    ) -> i32;
+
+    pub fn SherpaOnnxOfflineSpeakerDiarizationSetConfig(
+        sd: *const OfflineSpeakerDiarization,
+        config: *const OfflineSpeakerDiarizationConfig,
+    );
+
+    pub fn SherpaOnnxOfflineSpeakerDiarizationResultGetNumSpeakers(
+        r: *const OfflineSpeakerDiarizationResult,
+    ) -> i32;
+
+    pub fn SherpaOnnxOfflineSpeakerDiarizationResultGetNumSegments(
+        r: *const OfflineSpeakerDiarizationResult,
+    ) -> i32;
+
+    pub fn SherpaOnnxOfflineSpeakerDiarizationResultSortByStartTime(
+        r: *const OfflineSpeakerDiarizationResult,
+    ) -> *const OfflineSpeakerDiarizationSegment;
+
+    pub fn SherpaOnnxOfflineSpeakerDiarizationDestroySegment(
+        s: *const OfflineSpeakerDiarizationSegment,
+    );
+
+    pub fn SherpaOnnxOfflineSpeakerDiarizationProcess(
+        sd: *const OfflineSpeakerDiarization,
+        samples: *const c_float,
+        n: i32,
+    ) -> *const OfflineSpeakerDiarizationResult;
+
+    pub fn SherpaOnnxOfflineSpeakerDiarizationDestroyResult(
+        r: *const OfflineSpeakerDiarizationResult,
+    );
+}

sherpa-onnx/rust/sherpa-onnx/src/lib.rs

@@ -1,6 +1,7 @@
 mod audio_tagging;
 mod display;
 mod offline_asr;
+mod offline_speaker_diarization;
 mod offline_speech_denoiser;
 mod online_asr;
 mod online_punctuation;
@@ -15,6 +16,7 @@ mod wave;
 pub use audio_tagging::*;
 pub use display::*;
 pub use offline_asr::*;
+pub use offline_speaker_diarization::*;
 pub use offline_speech_denoiser::*;
 pub use online_asr::*;
 pub use online_punctuation::*;

sherpa-onnx/rust/sherpa-onnx/src/offline_speaker_diarization.rs

@@ -0,0 +1,231 @@
+use crate::{speaker_embedding::SpeakerEmbeddingExtractorConfig, utils::to_c_ptr};
+use sherpa_onnx_sys as sys;
+use std::ffi::CString;
+use std::slice;
+
+#[derive(Clone, Debug, Default)]
+pub struct OfflineSpeakerSegmentationPyannoteModelConfig {
+    pub model: Option<String>,
+}
+
+impl OfflineSpeakerSegmentationPyannoteModelConfig {
+    fn to_sys(
+        &self,
+        cstrings: &mut Vec<CString>,
+    ) -> sys::OfflineSpeakerSegmentationPyannoteModelConfig {
+        sys::OfflineSpeakerSegmentationPyannoteModelConfig {
+            model: to_c_ptr(&self.model, cstrings),
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct OfflineSpeakerSegmentationModelConfig {
+    pub pyannote: OfflineSpeakerSegmentationPyannoteModelConfig,
+    pub num_threads: i32,
+    pub debug: bool,
+    pub provider: Option<String>,
+}
+
+impl Default for OfflineSpeakerSegmentationModelConfig {
+    fn default() -> Self {
+        Self {
+            pyannote: Default::default(),
+            num_threads: 1,
+            debug: false,
+            provider: Some("cpu".to_string()),
+        }
+    }
+}
+
+impl OfflineSpeakerSegmentationModelConfig {
+    fn to_sys(&self, cstrings: &mut Vec<CString>) -> sys::OfflineSpeakerSegmentationModelConfig {
+        sys::OfflineSpeakerSegmentationModelConfig {
+            pyannote: self
+                .pyannote
+                .to_sys(cstrings),
+            num_threads: self.num_threads,
+            debug: self.debug as i32,
+            provider: to_c_ptr(&self.provider, cstrings),
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct FastClusteringConfig {
+    pub num_clusters: i32,
+    pub threshold: f32,
+}
+
+impl Default for FastClusteringConfig {
+    fn default() -> Self {
+        Self {
+            num_clusters: -1,
+            threshold: 0.5,
+        }
+    }
+}
+
+impl FastClusteringConfig {
+    fn to_sys(&self) -> sys::FastClusteringConfig {
+        sys::FastClusteringConfig {
+            num_clusters: self.num_clusters,
+            threshold: self.threshold,
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct OfflineSpeakerDiarizationConfig {
+    pub segmentation: OfflineSpeakerSegmentationModelConfig,
+    pub embedding: SpeakerEmbeddingExtractorConfig,
+    pub clustering: FastClusteringConfig,
+    pub min_duration_on: f32,
+    pub min_duration_off: f32,
+}
+
+impl Default for OfflineSpeakerDiarizationConfig {
+    fn default() -> Self {
+        Self {
+            segmentation: Default::default(),
+            embedding: Default::default(),
+            clustering: Default::default(),
+            min_duration_on: 0.3,
+            min_duration_off: 0.5,
+        }
+    }
+}
+
+impl OfflineSpeakerDiarizationConfig {
+    fn to_sys(&self, cstrings: &mut Vec<CString>) -> sys::OfflineSpeakerDiarizationConfig {
+        sys::OfflineSpeakerDiarizationConfig {
+            segmentation: self
+                .segmentation
+                .to_sys(cstrings),
+            embedding: self
+                .embedding
+                .to_sys(cstrings),
+            clustering: self
+                .clustering
+                .to_sys(),
+            min_duration_on: self.min_duration_on,
+            min_duration_off: self.min_duration_off,
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct OfflineSpeakerDiarizationSegment {
+    pub start: f32,
+    pub end: f32,
+    pub speaker: i32,
+}
+
+pub struct OfflineSpeakerDiarization {
+    ptr: *const sys::OfflineSpeakerDiarization,
+}
+
+unsafe impl Send for OfflineSpeakerDiarization {}
+
+impl OfflineSpeakerDiarization {
+    pub fn create(config: &OfflineSpeakerDiarizationConfig) -> Option<Self> {
+        let mut cstrings = Vec::new();
+        let sys_config = config.to_sys(&mut cstrings);
+        let ptr = unsafe { sys::SherpaOnnxCreateOfflineSpeakerDiarization(&sys_config) };
+        if ptr.is_null() {
+            None
+        } else {
+            Some(Self { ptr })
+        }
+    }
+
+    pub fn sample_rate(&self) -> i32 {
+        unsafe { sys::SherpaOnnxOfflineSpeakerDiarizationGetSampleRate(self.ptr) }
+    }
+
+    pub fn set_config(&self, config: &OfflineSpeakerDiarizationConfig) {
+        let mut cstrings = Vec::new();
+        let sys_config = config.to_sys(&mut cstrings);
+        unsafe { sys::SherpaOnnxOfflineSpeakerDiarizationSetConfig(self.ptr, &sys_config) }
+    }
+
+    pub fn process(&self, samples: &[f32]) -> Option<OfflineSpeakerDiarizationResult> {
+        let ptr = unsafe {
+            sys::SherpaOnnxOfflineSpeakerDiarizationProcess(
+                self.ptr,
+                samples.as_ptr(),
+                samples.len() as i32,
+            )
+        };
+        if ptr.is_null() {
+            None
+        } else {
+            Some(OfflineSpeakerDiarizationResult { ptr })
+        }
+    }
+}
+
+impl Drop for OfflineSpeakerDiarization {
+    fn drop(&mut self) {
+        unsafe {
+            if !self
+                .ptr
+                .is_null()
+            {
+                sys::SherpaOnnxDestroyOfflineSpeakerDiarization(self.ptr);
+            }
+        }
+    }
+}
+
+pub struct OfflineSpeakerDiarizationResult {
+    ptr: *const sys::OfflineSpeakerDiarizationResult,
+}
+
+impl OfflineSpeakerDiarizationResult {
+    pub fn num_speakers(&self) -> i32 {
+        unsafe { sys::SherpaOnnxOfflineSpeakerDiarizationResultGetNumSpeakers(self.ptr) }
+    }
+
+    pub fn num_segments(&self) -> i32 {
+        unsafe { sys::SherpaOnnxOfflineSpeakerDiarizationResultGetNumSegments(self.ptr) }
+    }
+
+    pub fn sort_by_start_time(&self) -> Vec<OfflineSpeakerDiarizationSegment> {
+        let n = self.num_segments();
+        if n <= 0 {
+            return Vec::new();
+        }
+
+        unsafe {
+            let p = sys::SherpaOnnxOfflineSpeakerDiarizationResultSortByStartTime(self.ptr);
+            if p.is_null() {
+                return Vec::new();
+            }
+
+            let segments = slice::from_raw_parts(p, n as usize)
+                .iter()
+                .map(|s| OfflineSpeakerDiarizationSegment {
+                    start: s.start,
+                    end: s.end,
+                    speaker: s.speaker,
+                })
+                .collect::<Vec<_>>();
+            sys::SherpaOnnxOfflineSpeakerDiarizationDestroySegment(p);
+            segments
+        }
+    }
+}
+
+impl Drop for OfflineSpeakerDiarizationResult {
+    fn drop(&mut self) {
+        unsafe {
+            if !self
+                .ptr
+                .is_null()
+            {
+                sys::SherpaOnnxOfflineSpeakerDiarizationDestroyResult(self.ptr);
+            }
+        }
+    }
+}

sherpa-onnx/rust/sherpa-onnx/src/speaker_embedding.rs

@@ -24,7 +24,7 @@ impl Default for SpeakerEmbeddingExtractorConfig {
 }
 
 impl SpeakerEmbeddingExtractorConfig {
-    fn to_sys(&self, cstrings: &mut Vec<CString>) -> sys::SpeakerEmbeddingExtractorConfig {
+    pub(crate) fn to_sys(&self, cstrings: &mut Vec<CString>) -> sys::SpeakerEmbeddingExtractorConfig {
         sys::SpeakerEmbeddingExtractorConfig {
             model: to_c_ptr(&self.model, cstrings),
             num_threads: self.num_threads,