Maven/cmd/mavwaked/vad.go

// Package main — mavwaked: always-on voice listening client.
//
// VAD state-machine: reads PCM 16k mono int16 windows, tracks speech↔silence
// transitions, and yields complete utterances as audio.Audio.
package main

import (
	"math"
	"time"

	"github.com/kami/maven/internal/audio"
)

// Defaults — same energy thresholds as mavsttd's gateReason (proven in
// silence_test.go) plus our own silence-hold timer.
const (
	defaultSampleRate = 16000
	defaultChannels   = 1
	defaultBits       = 16

	defaultFrameMs   = 30    // 480 samples — silero-vad window
	defaultSpeechMs  = 300   // min speech before accepting utterance
	defaultSilenceMs = 800   // silence hold before declaring end-of-utterance
	defaultMaxMs     = 10000 // cap single utterance at 10s
	defaultMinRMS    = 0.01  // RMS floor (same as mavsttd)

	// Barge-in thresholds. Only used when -barge-in is passed. The RMS is
	// x10000 like -min-rms, and sits an order of magnitude above the VAD's
	// own floor on purpose: with no acoustic echo canceller, a frame only
	// counts as "he is talking over her" if it is far louder than what the
	// speaker leaks back into the mic. 5 frames is 150ms — long enough that
	// a door or a cough does not cut her off mid-sentence.
	defaultBargeRMS    = 1200 // 0.12 normalised RMS
	defaultBargeFrames = 5
)

// frameSamples — samples per 30ms frame at 16kHz.
const frameSamples = defaultSampleRate * defaultFrameMs / 1000 // = 480

// SpeechState tracks whether the capture pipeline is listening or inside an
// utterance. Public so the caller can read the current state for logging.
type SpeechState int

const (
	StateSilence SpeechState = iota
	StateSpeech
)

// VAD is the energy-based voice activity detector.
//
// It consumes 30ms PCM frames, tracks the RMS energy floor adaptively, and
// signals when an utterance starts and ends. Same energy approach as
// mavsttd's gateReason — proven on real room audio in production.
type VAD struct {
	minRMS    float64
	speechMs  int // minimum ms of consecutive speech before triggering
	silenceMs int // ms of consecutive silence before ending utterance
	maxMs     int // absolute cap on utterance length

	// state
	state         SpeechState
	speechFrames  int    // consecutive speech frames during silence
	silenceFrames int    // consecutive silence frames during speech
	totalFrames   int    // frames since speech started (capped at max)
	utterance     []byte // raw PCM buffer for current utterance

	// adaptive noise floor — tracks recent silence RMS so the threshold
	// follows the room's ambient level. Initialised to minRMS; updated
	// on each silence frame.
	floorRMS float64

	// speech is silero-vad, or nil. When it is set the energy floor decides
	// nothing: the question becomes "is this speech" rather than "is this
	// loud", and the noise floor is not even tracked. Everything after that
	// answer — the speech hold, the silence hold, the length cap, the
	// buffer — is the same state machine either way, which is why the
	// detector goes here and not around this type.
	speech    *sileroVAD
	speechMin float64
}

// UseSilero swaps the energy threshold for the model. Passing nil is a
// no-op, so a caller that could not load the graph keeps a working VAD.
func (v *VAD) UseSilero(s *sileroVAD, threshold float64) {
	if s == nil {
		return
	}
	if threshold <= 0 {
		threshold = defaultSileroThreshold
	}
	v.speech = s
	v.speechMin = threshold
}

// isSpeech answers the one question the state machine asks of a frame.
func (v *VAD) isSpeech(frame []int16, rms float64) bool {
	if v.speech != nil {
		ok, _ := v.speech.Speech(frame, v.speechMin)
		return ok
	}
	return rms >= v.floorRMS
}

// NewVAD creates a VAD with the given thresholds. Zero values use defaults.
func NewVAD(minRMS, speechMs, silenceMs, maxMs int) *VAD {
	v := &VAD{
		minRMS:    defaultMinRMS,
		speechMs:  defaultSpeechMs,
		silenceMs: defaultSilenceMs,
		maxMs:     defaultMaxMs,
		floorRMS:  defaultMinRMS,
	}
	if minRMS > 0 {
		v.minRMS = float64(minRMS) / 10000.0
		v.floorRMS = v.minRMS
	}
	if speechMs > 0 {
		v.speechMs = speechMs
	}
	if silenceMs > 0 {
		v.silenceMs = silenceMs
	}
	if maxMs > 0 {
		v.maxMs = maxMs
	}
	return v
}

// FrameSamples returns the number of PCM int16 samples this VAD expects per
// feed call. Callers must chunk their stream accordingly.
func (v *VAD) FrameSamples() int { return frameSamples }

// State returns the current speech state.
func (v *VAD) State() SpeechState { return v.state }

// Feed processes one 30ms frame of PCM int16 LE samples. Returns:
//   - utterance: a complete utterance's PCM bytes, or nil if still accumulating
//   - state: the current speech state after processing this frame
//
// When an utterance is returned, the internal buffer resets and the caller
// should send the audio to the voice server before feeding more frames.
func (v *VAD) Feed(frame []int16) (_ audio.Audio, state SpeechState) {
	rms := frameRMS(frame)
	isSpeech := v.isSpeech(frame, rms)

	switch v.state {
	case StateSilence:
		if isSpeech {
			v.speechFrames++
			if v.speechFrames*frameSamples*1000/defaultSampleRate >= v.speechMs {
				// Transition to speech: start buffering utterance.
				v.state = StateSpeech
				v.silenceFrames = 0
				v.totalFrames = 0
				v.utterance = nil
				// Include the frames that built up to the threshold.
				v.utterance = append(v.utterance, pcmBytes(frame)...)
				return audio.Audio{}, StateSpeech
			}
		} else {
			// Silence — update noise floor (slow attack, fast decay).
			v.speechFrames = 0
			v.floorRMS = v.floorRMS*0.9 + rms*0.1
			if v.floorRMS < v.minRMS {
				v.floorRMS = v.minRMS
			}
		}
		return audio.Audio{}, StateSilence

	case StateSpeech:
		v.totalFrames++
		v.utterance = append(v.utterance, pcmBytes(frame)...)

		if v.totalFrames*frameSamples*1000/defaultSampleRate >= v.maxMs {
			// Max utterance length reached — force-end the utterance.
			utt := audio.Audio{
				Format: audio.PCM16kMono,
				Bytes:  v.utterance,
			}
			v.reset()
			return utt, StateSilence
		}

		if isSpeech {
			v.silenceFrames = 0
			return audio.Audio{}, StateSpeech
		}

		v.silenceFrames++
		if v.silenceFrames*frameSamples*1000/defaultSampleRate >= v.silenceMs {
			// Silence threshold reached — utterance complete.
			utt := audio.Audio{
				Format: audio.PCM16kMono,
				Bytes:  v.utterance,
			}
			v.reset()
			return utt, StateSilence
		}
		return audio.Audio{}, StateSpeech
	}

	return audio.Audio{}, StateSilence
}

// Reset clears the internal state (e.g. after a timeout or error).
func (v *VAD) Reset() { v.reset() }

func (v *VAD) reset() {
	if v.speech != nil {
		v.speech.Reset()
	}
	v.state = StateSilence
	v.speechFrames = 0
	v.silenceFrames = 0
	v.totalFrames = 0
	v.utterance = nil
}

// frameRMS computes the RMS energy of one int16 PCM frame.
func frameRMS(frame []int16) float64 {
	if len(frame) == 0 {
		return 0
	}
	var sum float64
	for _, s := range frame {
		f := float64(s) / 32768.0
		sum += f * f
	}
	return math.Sqrt(sum / float64(len(frame)))
}

// pcmBytes converts an int16 frame to raw PCM LE bytes.
func pcmBytes(frame []int16) []byte {
	b := make([]byte, len(frame)*2)
	for i, s := range frame {
		b[i*2] = byte(s)
		b[i*2+1] = byte(s >> 8)
	}
	return b
}

// PCMToF32 converts raw PCM LE bytes to float32 samples (normalized to [-1,1]).
// Used by tests and by the audio source to feed the VAD.
func PCMToF32(raw []byte) []float32 {
	n := len(raw) / 2
	out := make([]float32, n)
	for i := 0; i < n; i++ {
		s := int16(raw[i*2]) | int16(raw[i*2+1])<<8
		out[i] = float32(s) / 32768.0
	}
	return out
}

// PCMToI16 converts raw PCM LE bytes to int16 samples. Direct feed for VAD.
func PCMToI16(raw []byte) []int16 {
	n := len(raw) / 2
	out := make([]int16, n)
	for i := 0; i < n; i++ {
		out[i] = int16(raw[i*2]) | int16(raw[i*2+1])<<8
	}
	return out
}

// AudioDuration returns the duration in seconds of raw PCM 16k mono int16.
func AudioDuration(raw []byte) time.Duration {
	if len(raw) == 0 {
		return 0
	}
	samples := len(raw) / 2
	return time.Duration(samples) * time.Second / defaultSampleRate
}