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)
)

// 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
}

// 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 := rms >= v.floorRMS

	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() {
	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
}