Captioneer/src/captions/processor.go

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

import (
	"strings"

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

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

	fixedSamples   []float32
	previewSamples []float32
	previewActive  bool
	previewStart   int
	nextPreviewAt  int
	totalSamples   int
}

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

func (p *Processor) Accept(samples []float32) error {
	if len(samples) == 0 {
		return nil
	}
	if p.settings.Mode == config.ModeFixed {
		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.emitFinal(p.totalSamples, p.totalSamples+len(p.fixedSamples), p.fixedSamples)
		}
		return nil
	}
	p.transcriber.FlushVAD()
	return p.drainVAD()
}

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.emitFinal(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)

	if 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...)
		if len(p.previewSamples) >= p.nextPreviewAt {
			if err := p.emitProvisional(); err != nil {
				return err
			}
			p.nextPreviewAt += audio.DurationSamples(p.settings.ChunkDuration)
		}
	}
	return p.drainVAD()
}

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.previewSamples = nil
		p.previewActive = false
		p.previewStart = 0
		p.nextPreviewAt = audio.DurationSamples(p.settings.ChunkDuration)
	}
}

func (p *Processor) emitProvisional() error {
	text := strings.TrimSpace(p.transcriber.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.transcriber.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),
	})
}