Captioneer/src/captions/processor.go

// Package captions turns audio packets into fixed or speech-aware captions.
package captions

import (
	"fmt"
	"io"

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

type Processor struct {
	settings  config.Settings
	resources *models.Resources
	output    io.Writer

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

func New(settings config.Settings, resources *models.Resources, output io.Writer) *Processor {
	return &Processor{
		settings:      settings,
		resources:     resources,
		output:        output,
		nextPreviewAt: audio.DurationSamples(settings.ChunkDuration),
	}
}

func (p *Processor) Accept(samples []float32) {
	if len(samples) == 0 {
		return
	}
	if p.settings.Mode == config.ModeFixed {
		p.acceptFixed(samples)
		return
	}
	p.acceptVAD(samples)
}

func (p *Processor) Flush() {
	if p.settings.Mode == config.ModeFixed {
		if len(p.fixedSamples) > 0 {
			p.emitCommitted(p.totalSamples, p.totalSamples+len(p.fixedSamples), p.fixedSamples)
		}
		return
	}
	p.resources.FlushVAD()
	p.drainVAD()
	p.clearPreview()
}

func (p *Processor) acceptFixed(samples []float32) {
	p.fixedSamples = append(p.fixedSamples, samples...)
	chunkSize := audio.DurationSamples(p.settings.ChunkDuration)
	for len(p.fixedSamples) >= chunkSize {
		p.emitCommitted(p.totalSamples, p.totalSamples+chunkSize, p.fixedSamples[:chunkSize])
		p.fixedSamples = p.fixedSamples[chunkSize:]
		p.totalSamples += chunkSize
	}
}

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

	if p.previewActive || audio.RMSDBFS(samples) >= p.settings.PreviewThresholdDB {
		p.previewActive = true
		p.previewSamples = append(p.previewSamples, samples...)
		if len(p.previewSamples) >= p.nextPreviewAt {
			p.emitPreview(p.previewSamples)
			p.nextPreviewAt += audio.DurationSamples(p.settings.ChunkDuration)
		}
	}
	p.drainVAD()
}

func (p *Processor) drainVAD() {
	for {
		segment, ok := p.resources.NextSpeechSegment()
		if !ok {
			return
		}
		p.clearPreview()
		p.emitCommitted(segment.Start, segment.Start+len(segment.Samples), segment.Samples)
		p.previewSamples = nil
		p.previewActive = false
		p.nextPreviewAt = audio.DurationSamples(p.settings.ChunkDuration)
	}
}

func (p *Processor) emitPreview(samples []float32) {
	text := p.resources.Decode(samples)
	if text == "" {
		return
	}
	fmt.Fprintf(p.output, "\r\033[2K… %s", text)
	p.previewVisible = true
}

func (p *Processor) emitCommitted(start, end int, samples []float32) {
	text := p.resources.Decode(samples)
	if text == "" {
		return
	}
	p.clearPreview()
	fmt.Fprintf(p.output, "[%s-%s] %s\n", audio.FormatTime(start), audio.FormatTime(end), text)
}

func (p *Processor) clearPreview() {
	if p.previewVisible {
		fmt.Fprint(p.output, "\r\033[2K")
		p.previewVisible = false
	}
}