reference/cpp/vad_iterator_onnx.cpp

#include <cmath>
#include <stdexcept>
#include <unordered_map>
#include <vector>

#include "vad_iterator_onnx.h"

VadIteratorOnnx::VadIteratorOnnx(float threshold,
                         int sampling_rate,
                         int min_silence_duration_ms,
                         float max_speech_duration_s,
                         int speech_pad_ms)
    : threshold_(threshold),
      sampling_rate_(sampling_rate),
      min_silence_samples_(sampling_rate_ * min_silence_duration_ms / 1000.0),
      speech_pad_samples_(sampling_rate_ * speech_pad_ms / 1000.0),
      triggered_(false),
      temp_end_(0),
      current_sample_(0),
      start_(0) {

    if (sampling_rate_ != 8000 && sampling_rate_ != 16000) {
        throw std::invalid_argument("Only support sampling rates of 8000 or 16000");
    }

    model_ = std::make_unique<OnnxVadWrapper>("path/to/vad.onnx", true); // 可配置路径
}

VadIteratorOnnx::~VadIteratorOnnx() = default;

void VadIteratorOnnx::reset_states() {
    model_->reset_states();
    triggered_ = false;
    temp_end_ = 0;
    current_sample_ = 0;
    start_ = 0;
    buffer_.clear();
}

std::unordered_map<std::string, double> 
VadIteratorOnnx::operator()(const std::vector<float>& x, bool return_seconds) {
    std::unordered_map<std::string, double> result;

    int window_size_samples = (sampling_rate_ == 16000) ? 512 : 256;
    
    // 将新音频追加到缓存中
    buffer_.insert(buffer_.end(), x.begin(), x.end());
    while (buffer_.size() > 0) {
        std::unordered_map<std::string, double> tmp;
        std::vector<float> chunk(buffer_.begin(), buffer_.begin() + std::min(static_cast<int>(x.size()), window_size_samples));
        // 补零到固定长度
        if (chunk.size() < static_cast<size_t>(window_size_samples)) {
            chunk.resize(window_size_samples, 0.0f);
        }

        current_sample_ += window_size_samples;

        // 推理得到语音概率
        auto [output, _] = (*model_)(chunk, sampling_rate_);
        float speech_prob = output[0];

        if (speech_prob >= threshold_ && temp_end_ > 0) {
            temp_end_ = 0;
        }

        if (speech_prob >= threshold_ && !triggered_) {
            triggered_ = true;
            start_ = std::max(0.0, current_sample_ - speech_pad_samples_ - window_size_samples);
            tmp["start"] = return_seconds ? start_ / sampling_rate_ : start_;
        }

        if (speech_prob < (threshold_ - 0.15) && triggered_) {
            if (temp_end_ == 0) {
                temp_end_ = current_sample_;
            }

            if (current_sample_ - temp_end_ >= min_silence_samples_) {
                double speech_end = temp_end_ + speech_pad_samples_ - window_size_samples;
                tmp["end"] = return_seconds ? speech_end / sampling_rate_ : speech_end;
                temp_end_ = 0;
                triggered_ = false;
            }
        }

        // 移除已处理的数据
        std::vector<float>(buffer_.begin() + window_size_samples, buffer_.end()).swap(buffer_);
        
        if (result.empty()) {
            result = tmp;
        } else if (!tmp.empty()) {
            // 如果当前结果有 'end'，更新最终 end
            if (tmp.find("end") != tmp.end()) {
                result["end"] = tmp["end"];
            }

            // 如果有新的 start，但前一个有 end，则合并成连续语音段
            if (tmp.find("start") != tmp.end() && result.find("end") != result.end()) {
                result.erase("end");
            }
        }
    }

    return result;
}