Captioneer/src/captions/processor.go

// Package captions turns audio packets into presentation-neutral caption events.
package captions

import (
	"strings"
	"time"

	"tea.chunkbyte.com/kato/captioneer/src/audio"
	"tea.chunkbyte.com/kato/captioneer/src/config"
)

// Provisional recognition uses only a recent window. Re-decoding an entire
// long utterance every refresh grows increasingly expensive, while the UI only
// presents the latest provisional lines. The final event still decodes the
// complete VAD segment.
const maxPreviewDuration = 10 * time.Second

type Processor struct {
	settings    config.Settings
	transcriber Transcriber
	emit        func(Event) error

	fixedSamples    []float32
	previewSamples  []float32
	previewActive   bool
	previewStart    int
	nextPreviewAt   int
	previewReceived int
	totalSamples    int
	activity        Activity
	speechKnown     bool
	speechActive    bool
}

type Activity struct {
	SpeechChanged      func(bool)
	RecognitionChanged func(bool)
}

func New(settings config.Settings, transcriber Transcriber, emit func(Event) error) *Processor {
	return NewWithActivity(settings, transcriber, emit, Activity{})
}

func NewWithActivity(settings config.Settings, transcriber Transcriber, emit func(Event) error, activity Activity) *Processor {
	return &Processor{
		settings:      settings,
		transcriber:   transcriber,
		emit:          emit,
		nextPreviewAt: audio.DurationSamples(settings.ChunkDuration),
		activity:      activity,
	}
}

func (p *Processor) Accept(samples []float32) error {
	if len(samples) == 0 {
		return nil
	}
	if p.settings.Mode == config.ModeFixed {
		p.setSpeechActive(audio.RMSDBFS(samples) >= p.settings.PreviewThresholdDB)
		return p.acceptFixed(samples)
	}
	return p.acceptVAD(samples)
}

func (p *Processor) Flush() error {
	if p.settings.Mode == config.ModeFixed {
		if len(p.fixedSamples) > 0 {
			return p.emitFixedFinal(p.totalSamples, p.totalSamples+len(p.fixedSamples), p.fixedSamples)
		}
		return nil
	}
	p.transcriber.FlushVAD()
	if err := p.drainVAD(); err != nil {
		return err
	}
	if p.previewActive {
		p.resetPreview()
		return p.emit(Event{Kind: Hide})
	}
	return nil
}

func (p *Processor) acceptFixed(samples []float32) error {
	p.fixedSamples = append(p.fixedSamples, samples...)
	chunkSize := audio.DurationSamples(p.settings.ChunkDuration)
	for len(p.fixedSamples) >= chunkSize {
		if err := p.emitFixedFinal(p.totalSamples, p.totalSamples+chunkSize, p.fixedSamples[:chunkSize]); err != nil {
			return err
		}
		p.fixedSamples = p.fixedSamples[chunkSize:]
		p.totalSamples += chunkSize
	}
	return nil
}

func (p *Processor) acceptVAD(samples []float32) error {
	p.totalSamples += len(samples)
	p.transcriber.AcceptVAD(samples)

	speechActive := p.transcriber.SpeechActive()
	p.setSpeechActive(speechActive)
	if speechActive && (p.previewActive || audio.RMSDBFS(samples) >= p.settings.PreviewThresholdDB) {
		if !p.previewActive {
			p.previewActive = true
			p.previewStart = p.totalSamples - len(samples)
		}
		p.previewSamples = append(p.previewSamples, samples...)
		p.previewReceived += len(samples)
		p.trimPreviewWindow()
		if p.previewReceived >= p.nextPreviewAt {
			if err := p.emitProvisional(); err != nil {
				return err
			}
			p.nextPreviewAt += audio.DurationSamples(p.settings.ChunkDuration)
		}
	}
	if err := p.drainVAD(); err != nil {
		return err
	}
	// A short VAD false-positive can become active without meeting the minimum
	// speech duration, so no final segment is produced. Clear its provisional
	// text instead of leaving a stale caption on screen.
	if !speechActive && p.previewActive {
		p.resetPreview()
		return p.emit(Event{Kind: Hide})
	}
	return nil
}

func (p *Processor) drainVAD() error {
	for {
		start, samples, ok := p.transcriber.NextSpeechSegment()
		if !ok {
			return nil
		}
		if err := p.emitFinal(start, start+len(samples), samples); err != nil {
			return err
		}
		p.resetPreview()
	}
}

func (p *Processor) trimPreviewWindow() {
	maxSamples := audio.DurationSamples(maxPreviewDuration)
	if len(p.previewSamples) <= maxSamples {
		return
	}
	dropped := len(p.previewSamples) - maxSamples
	p.previewSamples = p.previewSamples[dropped:]
	p.previewStart += dropped
}

func (p *Processor) resetPreview() {
	p.previewSamples = nil
	p.previewActive = false
	p.previewStart = 0
	p.previewReceived = 0
	p.nextPreviewAt = audio.DurationSamples(p.settings.ChunkDuration)
}

func (p *Processor) emitProvisional() error {
	text := strings.TrimSpace(p.decode(p.previewSamples))
	if text == "" {
		return nil
	}
	return p.emit(Event{
		Kind:      Provisional,
		Text:      text,
		StartedAt: audio.SamplesDuration(p.previewStart),
		EndedAt:   audio.SamplesDuration(p.previewStart + len(p.previewSamples)),
	})
}

func (p *Processor) emitFinal(start, end int, samples []float32) error {
	text := strings.TrimSpace(p.decode(samples))
	if text == "" {
		return p.emit(Event{Kind: Hide})
	}
	return p.emit(Event{
		Kind:      Final,
		Text:      text,
		StartedAt: audio.SamplesDuration(start),
		EndedAt:   audio.SamplesDuration(end),
	})
}

func (p *Processor) decode(samples []float32) string {
	if p.activity.RecognitionChanged != nil {
		p.activity.RecognitionChanged(true)
		defer p.activity.RecognitionChanged(false)
	}
	return p.transcriber.Decode(samples)
}

func (p *Processor) setSpeechActive(active bool) {
	if p.speechKnown && p.speechActive == active {
		return
	}
	p.speechKnown = true
	p.speechActive = active
	if p.activity.SpeechChanged != nil {
		p.activity.SpeechChanged(active)
	}
}

func (p *Processor) emitFixedFinal(start, end int, samples []float32) error {
	if audio.RMSDBFS(samples) < p.settings.PreviewThresholdDB {
		return p.emit(Event{Kind: Hide})
	}
	return p.emitFinal(start, end, samples)
}