voxudio 0.7.1

A real-time audio processing library with ONNX runtime support
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

use {
    super::get_session_builder,
    crate::OperationError,
    ndarray::Array2,
    ort::{
        inputs,
        session::{RunOptions, Session},
        value::TensorRef,
    },
    std::{array::from_fn, path::Path},
};

/// 说话人嵌入提取器，用于从音频中提取说话人特征（音色）向量
///
/// # 示例
/// ```
/// use voxudio::{SpeakerEmbeddingExtractor, load_audio};
/// #[tokio::main]
/// async fn main() -> anyhow::Result<()> {
/// let mut see = SpeakerEmbeddingExtractor::new("../checkpoint/speaker_embedding_extractor.onnx")?;
/// let (audio_data, channels) = load_audio::<22050, f32, _>("../asset/hello_in_cn.mp3", true).await?;
/// let embeddings = see.extract(&audio_data, channels).await?;
/// Ok(())
/// }
/// ```
///
/// # 错误
/// - 模型加载失败时返回 `OperationError`
/// - 音频处理失败时返回 `OperationError`
///
/// # 注意事项
/// - 只能处理22050采样率的音频
/// - 输入音频数据应为单声道或双声道PCM格式
/// - 每个声道会输出一个256维的特征向量
pub struct SpeakerEmbeddingExtractor {
    model: Session,
}

impl SpeakerEmbeddingExtractor {
    /// 创建一个新的说话人嵌入提取器实例
    ///
    /// # 参数
    /// - `model_path`: ONNX模型文件路径
    ///
    /// # 返回值
    /// - `Result<Self, OperationError>`: 成功返回提取器实例，失败返回错误
    ///
    /// # 示例
    /// ```
    /// use voxudio::SpeakerEmbeddingExtractor;
    /// fn main() -> anyhow::Result<()> {
    /// let extractor = SpeakerEmbeddingExtractor::new("../checkpoint/speaker_embedding_extractor.onnx")?;
    ///
    /// Ok(())
    /// }
    /// ```
    pub fn new<P>(model_path: P) -> Result<Self, OperationError>
    where
        P: AsRef<Path>,
    {
        let model = get_session_builder()?.commit_from_file(model_path)?;

        Ok(Self { model })
    }

    /// 从音频数据中提取说话人嵌入特征
    ///
    /// # 参数
    /// - `audio`: 音频数据，采样率必须为22050Hz
    ///
    /// # 返回值
    /// - `Result<Vec<[f32; 256]>, OperationError>`: 成功返回每个声道的256维特征向量，失败返回错误
    ///
    /// # 示例
    /// ```
    /// use voxudio::SpeakerEmbeddingExtractor;
    /// #[tokio::main]
    /// async fn main() -> anyhow::Result<()> {
    ///     let mut see = SpeakerEmbeddingExtractor::new("../checkpoint/speaker_embedding_extractor.onnx")?;
    ///     let embeddings = see.extract(&[0.0; 22050], 1).await?;
    ///     Ok(())
    /// }
    /// ```
    pub async fn extract(
        &mut self,
        audio: &[f32],
        channels: usize,
    ) -> Result<Vec<[f32; 256]>, OperationError> {
        let len = audio.len();
        let x = Array2::from_shape_vec((len / channels, channels), audio.to_vec())?;
        let options = RunOptions::new()?;
        let outputs = self
            .model
            .run_async(
                inputs![
                    "audio" => TensorRef::from_array_view(&x)?
                ],
                &options,
            )?
            .await?;
        let se = outputs["se"].try_extract_array::<f32>()?;

        let mut out = Vec::with_capacity(channels);
        // 使用数组推导式更简洁地填充结果
        for i in 0..channels {
            let row: [f32; 256] = from_fn(|j| se[[i, j]]);
            out.push(row);
        }

        Ok(out)
    }
}