record-asr/internal/audio/decode.go

// Package audio decodes uploaded audio files into the 16 kHz mono float32
// sample slice that sherpa-onnx expects.
//
// Pure-Go beep decoders handle .wav/.mp3/.flac/.ogg. Everything else
// (notably .m4a/.aac — there is no robust pure-Go AAC decoder) is decoded by
// shelling out to ffmpeg, which converts to 16 kHz mono f32le PCM on stdout.
package audio

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"math"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/gopxl/beep/v2"
	"github.com/gopxl/beep/v2/flac"
	"github.com/gopxl/beep/v2/mp3"
	"github.com/gopxl/beep/v2/vorbis"
	"github.com/gopxl/beep/v2/wav"
)

// TargetSampleRate is the rate sherpa-onnx models were trained on.
const TargetSampleRate = 16000

// Package-level decoder configuration set once by the server bootstrap.
var (
	ffmpegPath  = "ffmpeg"
	decTimeout  = 30 * time.Minute
)

// Configure sets the ffmpeg binary path and a per-file decode timeout. Call
// once at startup from the config. Safe to call multiple times.
func Configure(ffmpeg string, timeout time.Duration) {
	if strings.TrimSpace(ffmpeg) != "" {
		ffmpegPath = ffmpeg
	}
	if timeout > 0 {
		decTimeout = timeout
	}
}

// beepExts are the formats beep decodes natively (pure Go, no external tool).
var beepExts = map[string]bool{
	".wav": true, ".mp3": true, ".flac": true, ".ogg": true, ".oga": true,
}

// Decode loads the file at path and returns mono float32 samples resampled to
// TargetSampleRate. Beep handles its native formats; ffmpeg handles the rest
// (m4a, aac, opus, wma, webm, aiff, …). If a beep-native file fails to decode,
// it transparently falls back to ffmpeg so a quirky/mis-named file still works.
func Decode(path string) (samples []float32, sampleRate int, err error) {
	ext := strings.ToLower(filepath.Ext(path))

	if beepExts[ext] {
		if s, format, e := tryBeep(path); e == nil {
			out, e2 := finalizeBeep(s, format)
			if e2 == nil {
				return out, TargetSampleRate, nil
			}
			// beep read failed mid-stream; fall through to ffmpeg.
		}
		// beep decode failed entirely; fall through to ffmpeg.
	}
	return decodeFFmpeg(path)
}

// tryBeep opens the file and dispatches to the matching beep decoder.
func tryBeep(path string) (beep.StreamSeekCloser, beep.Format, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, beep.Format{}, fmt.Errorf("open audio: %w", err)
	}
	switch strings.ToLower(filepath.Ext(path)) {
	case ".wav":
		return wav.Decode(f)
	case ".mp3":
		return mp3.Decode(f)
	case ".flac":
		return flac.Decode(f)
	case ".ogg", ".oga":
		return vorbis.Decode(f)
	}
	// Caller should not reach here for non-beep exts.
	f.Close()
	return nil, beep.Format{}, fmt.Errorf("not a beep-supported extension")
}

// finalizeBeep drains the decoder into a mono float64 buffer at the native rate
// and resamples to 16 kHz mono float32.
func finalizeBeep(s beep.StreamSeekCloser, format beep.Format) ([]float32, error) {
	defer s.Close()
	const bufLen = 8192
	buf := make([][2]float64, bufLen)
	mono := make([]float64, 0, 1<<20)
	for {
		n, ok := s.Stream(buf)
		if n > 0 {
			for i := 0; i < n; i++ {
				mono = append(mono, (buf[i][0]+buf[i][1])/2.0)
			}
		}
		if !ok {
			break
		}
	}
	if err := s.Err(); err != nil {
		return nil, fmt.Errorf("read audio: %w", err)
	}
	if len(mono) == 0 {
		return nil, fmt.Errorf("audio file has no samples")
	}
	return resampleMono(mono, int(format.SampleRate), TargetSampleRate), nil
}

// decodeFFmpeg runs `ffmpeg -i path -f f32le -ar 16000 -ac 1 -` and reads the
// raw little-endian float32 PCM from stdout. This is the universal fallback
// for formats beep cannot decode (m4a/aac/opus/wma/…).
func decodeFFmpeg(path string) ([]float32, int, error) {
	if _, err := exec.LookPath(ffmpegPath); err != nil {
		return nil, 0, fmt.Errorf("ffmpeg not found (%s): install ffmpeg or set ASR_FFMPEG to decode this format", ffmpegPath)
	}
	ctx, cancel := context.WithTimeout(context.Background(), decTimeout)
	defer cancel()

	var stderr bytes.Buffer
	// -vn drops the video stream entirely: for video inputs (mp4/mkv/…) this
	// skips video decoding so only the audio track is extracted; for
	// audio-only inputs it's a harmless no-op.
	cmd := exec.CommandContext(ctx, ffmpegPath,
		"-nostdin", "-hide_banner", "-loglevel", "error",
		"-i", path,
		"-vn",
		"-f", "f32le", "-acodec", "pcm_f32le",
		"-ar", "16000", "-ac", "1",
		"-",
	)
	var stdout bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		if ctx.Err() == context.DeadlineExceeded {
			return nil, 0, fmt.Errorf("ffmpeg decode timed out after %s", decTimeout)
		}
		msg := strings.TrimSpace(stderr.String())
		if msg == "" {
			msg = err.Error()
		}
		return nil, 0, fmt.Errorf("ffmpeg decode failed: %s", msg)
	}
	raw := stdout.Bytes()
	if len(raw) == 0 {
		return nil, 0, fmt.Errorf("ffmpeg produced no audio (empty or no audio stream in %s)", filepath.Base(path))
	}
	if len(raw)%4 != 0 {
		raw = raw[:len(raw)-len(raw)%4] // drop trailing partial sample
	}
	out := make([]float32, len(raw)/4)
	for i := 0; i < len(out); i++ {
		out[i] = math.Float32frombits(binary.LittleEndian.Uint32(raw[i*4 : i*4+4]))
	}
	return out, TargetSampleRate, nil
}

// resampleMono converts a mono float64 slice from srcRate to dstRate using
// linear interpolation. Adequate for speech recognition at 16 kHz.
func resampleMono(in []float64, srcRate, dstRate int) []float32 {
	if srcRate == dstRate {
		out := make([]float32, len(in))
		for i, v := range in {
			out[i] = float32(v)
		}
		return out
	}
	if srcRate <= 0 || dstRate <= 0 || len(in) < 2 {
		out := make([]float32, len(in))
		for i, v := range in {
			out[i] = float32(v)
		}
		return out
	}
	ratio := float64(dstRate) / float64(srcRate)
	outLen := int(float64(len(in)-1) * ratio)
	if outLen < 1 {
		outLen = 1
	}
	out := make([]float32, outLen)
	for i := 0; i < outLen; i++ {
		pos := float64(i) / ratio
		idx := int(pos)
		if idx >= len(in)-1 {
			out[i] = float32(in[len(in)-1])
			continue
		}
		frac := pos - float64(idx)
		v := in[idx]*(1-frac) + in[idx+1]*frac
		out[i] = float32(v)
	}
	return out
}

// acceptedExts is the allowlist the upload handler rejects against. Beep-native
// formats decode without ffmpeg; the rest need ffmpeg on PATH.
var acceptedExts = map[string]bool{
	".wav": true, ".mp3": true, ".flac": true, ".ogg": true, ".oga": true,
	".m4a": true, ".m4b": true, ".aac": true, ".opus": true,
	".wma": true, ".webm": true, ".aiff": true, ".aif": true,
	".amr": true, ".au": true, ".snd": true, ".gsm": true,
	// Video containers — the audio track is extracted via ffmpeg (-vn drops video).
	".mp4": true, ".m4v": true, ".mkv": true, ".mov": true, ".avi": true,
	".ts": true, ".m2ts": true, ".mts": true, ".wmv": true, ".flv": true,
	".3gp": true, ".3g2": true, ".mpg": true, ".mpeg": true, ".vob": true,
}

// SupportedExtensions lists the file extensions Decode accepts.
func SupportedExtensions() string {
	return ".wav, .mp3, .flac, .ogg, .m4a, .aac, .opus, .aiff, + 视频 .mp4/.mkv/.mov/.avi/.ts… (ffmpeg extracts audio)"
}

// IsSupported reports whether the file at path has a decodable extension.
func IsSupported(path string) bool {
	return acceptedExts[strings.ToLower(filepath.Ext(path))]
}