diff --git a/.github/scripts/test-rust.sh b/.github/scripts/test-rust.sh index 3827d60f9d..376c427c88 100755 --- a/.github/scripts/test-rust.sh +++ b/.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 diff --git a/rust-api-examples/README.md b/rust-api-examples/README.md index 9f2192279f..9f3c0919c9 100644 --- a/rust-api-examples/README.md +++ b/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 +``` diff --git a/rust-api-examples/examples/offline_speaker_diarization.rs b/rust-api-examples/examples/offline_speaker_diarization.rs new file mode 100644 index 0000000000..a34a1ec80e --- /dev/null +++ b/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); + } +} diff --git a/rust-api-examples/run-offline-speaker-diarization.sh b/rust-api-examples/run-offline-speaker-diarization.sh new file mode 100755 index 0000000000..cf57f3b074 --- /dev/null +++ b/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 diff --git a/sherpa-onnx/rust/sherpa-onnx-sys/src/lib.rs b/sherpa-onnx/rust/sherpa-onnx-sys/src/lib.rs index f165aa77b8..73af38fd16 100644 --- a/sherpa-onnx/rust/sherpa-onnx-sys/src/lib.rs +++ b/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::*; diff --git a/sherpa-onnx/rust/sherpa-onnx-sys/src/offline_speaker_diarization.rs b/sherpa-onnx/rust/sherpa-onnx-sys/src/offline_speaker_diarization.rs new file mode 100644 index 0000000000..00232e0a98 --- /dev/null +++ b/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, + ); +} diff --git a/sherpa-onnx/rust/sherpa-onnx/src/lib.rs b/sherpa-onnx/rust/sherpa-onnx/src/lib.rs index 275bc3f34d..5fd9be9e11 100644 --- a/sherpa-onnx/rust/sherpa-onnx/src/lib.rs +++ b/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::*; diff --git a/sherpa-onnx/rust/sherpa-onnx/src/offline_speaker_diarization.rs b/sherpa-onnx/rust/sherpa-onnx/src/offline_speaker_diarization.rs new file mode 100644 index 0000000000..d8c0458739 --- /dev/null +++ b/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, +} + +impl OfflineSpeakerSegmentationPyannoteModelConfig { + fn to_sys( + &self, + cstrings: &mut Vec, + ) -> 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, +} + +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) -> 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) -> 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 { + 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 { + 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 { + 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::>(); + sys::SherpaOnnxOfflineSpeakerDiarizationDestroySegment(p); + segments + } + } +} + +impl Drop for OfflineSpeakerDiarizationResult { + fn drop(&mut self) { + unsafe { + if !self + .ptr + .is_null() + { + sys::SherpaOnnxOfflineSpeakerDiarizationDestroyResult(self.ptr); + } + } + } +} diff --git a/sherpa-onnx/rust/sherpa-onnx/src/speaker_embedding.rs b/sherpa-onnx/rust/sherpa-onnx/src/speaker_embedding.rs index 79a33df209..2f6080b11c 100644 --- a/sherpa-onnx/rust/sherpa-onnx/src/speaker_embedding.rs +++ b/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) -> sys::SpeakerEmbeddingExtractorConfig { + pub(crate) fn to_sys(&self, cstrings: &mut Vec) -> sys::SpeakerEmbeddingExtractorConfig { sys::SpeakerEmbeddingExtractorConfig { model: to_c_ptr(&self.model, cstrings), num_threads: self.num_threads,