Pad audio before WhisperX transcription

internal/transcription/client.go

@@ -8,14 +8,18 @@ import (
 	"io"
 	"mime/multipart"
 	"net/http"
+	"os"
+	"os/exec"
 	"path/filepath"
 	"strings"
 	"time"
 )
 
 type Client struct {
-	baseURL string
-	http    *http.Client
+	baseURL     string
+	http        *http.Client
+	ffmpegPath  string
+	leadSilence time.Duration
 }
 
 type Input struct {
@@ -49,7 +53,7 @@ type whisperResponse struct {
 	AlignError   *string   `json:"align_error,omitempty"`
 }
 
-func New(baseURL string, timeout time.Duration) *Client {
+func New(baseURL string, timeout time.Duration, ffmpegPath string, leadSilence time.Duration) *Client {
 	baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
 	if baseURL == "" {
 		return nil
@@ -57,9 +61,21 @@ func New(baseURL string, timeout time.Duration) *Client {
 	if timeout <= 0 {
 		timeout = 10 * time.Minute
 	}
+	if leadSilence < 0 {
+		leadSilence = 0
+	}
+	if leadSilence > 5*time.Second {
+		leadSilence = 5 * time.Second
+	}
+	ffmpegPath = strings.TrimSpace(ffmpegPath)
+	if ffmpegPath == "" {
+		ffmpegPath = "ffmpeg"
+	}
 	return &Client{
-		baseURL: baseURL,
-		http:    &http.Client{Timeout: timeout},
+		baseURL:     baseURL,
+		http:        &http.Client{Timeout: timeout},
+		ffmpegPath:  ffmpegPath,
+		leadSilence: leadSilence,
 	}
 }
 
@@ -74,13 +90,20 @@ func (c *Client) Transcribe(ctx context.Context, in Input) (*Result, error) {
 	if err != nil {
 		return nil, err
 	}
+	if c.leadSilence > 0 {
+		audio, filename, err = c.addLeadSilence(ctx, audio, filename)
+		if err != nil {
+			return nil, err
+		}
+	}
 	resp, duration, err := c.transcribeAudio(ctx, audio, filename, in)
 	if err != nil {
 		return nil, err
 	}
+	segments := adjustLeadSilence(resp.Segments, c.leadSilence)
 	return &Result{
 		Language:     resp.Language,
-		Segments:     resp.Segments,
+		Segments:     segments,
 		DiarizeError: resp.DiarizeError,
 		AlignError:   resp.AlignError,
 		DurationMS:   duration.Milliseconds(),
@@ -115,6 +138,90 @@ func (c *Client) downloadAudio(ctx context.Context, in Input) ([]byte, string, e
 	return audio, filename, nil
 }
 
+func (c *Client) addLeadSilence(ctx context.Context, audio []byte, filename string) ([]byte, string, error) {
+	tmpDir, err := os.MkdirTemp("", "ai-transcribe-*")
+	if err != nil {
+		return nil, "", fmt.Errorf("prepare audio temp dir: %w", err)
+	}
+	defer os.RemoveAll(tmpDir)
+
+	inputPath := filepath.Join(tmpDir, "input"+safeExt(filename))
+	outputPath := filepath.Join(tmpDir, "padded.mp3")
+	if err := os.WriteFile(inputPath, audio, 0o600); err != nil {
+		return nil, "", fmt.Errorf("write audio temp file: %w", err)
+	}
+	delayMS := int(c.leadSilence.Milliseconds())
+	if delayMS <= 0 {
+		return audio, filename, nil
+	}
+	cmd := exec.CommandContext(ctx, c.ffmpegPath,
+		"-nostdin", "-y",
+		"-i", inputPath,
+		"-af", fmt.Sprintf("adelay=%d:all=1", delayMS),
+		"-codec:a", "libmp3lame",
+		"-qscale:a", "5",
+		outputPath,
+	)
+	out, err := cmd.CombinedOutput()
+	if err != nil {
+		return nil, "", fmt.Errorf("ffmpeg lead silence: %w (%s)", err, trimOutput(out))
+	}
+	padded, err := os.ReadFile(outputPath)
+	if err != nil {
+		return nil, "", fmt.Errorf("read padded audio: %w", err)
+	}
+	if len(padded) == 0 {
+		return nil, "", fmt.Errorf("padded audio is empty")
+	}
+	base := strings.TrimSuffix(filepath.Base(filename), filepath.Ext(filename))
+	if base == "" || base == "." || base == "/" {
+		base = "audio"
+	}
+	return padded, base + "-padded.mp3", nil
+}
+
+func safeExt(filename string) string {
+	ext := strings.ToLower(filepath.Ext(filename))
+	switch ext {
+	case ".mp3", ".wav", ".m4a", ".ogg", ".opus", ".webm":
+		return ext
+	default:
+		return ".audio"
+	}
+}
+
+func trimOutput(out []byte) string {
+	s := strings.TrimSpace(string(out))
+	if len(s) > 600 {
+		return s[:600]
+	}
+	return s
+}
+
+func adjustLeadSilence(segments []Segment, silence time.Duration) []Segment {
+	if len(segments) == 0 || silence <= 0 {
+		return segments
+	}
+	shift := silence.Seconds()
+	out := make([]Segment, 0, len(segments))
+	for _, segment := range segments {
+		segment.Start = clampTime(segment.Start - shift)
+		segment.End = clampTime(segment.End - shift)
+		if segment.End < segment.Start {
+			segment.End = segment.Start
+		}
+		out = append(out, segment)
+	}
+	return out
+}
+
+func clampTime(v float64) float64 {
+	if v < 0 {
+		return 0
+	}
+	return v
+}
+
 func (c *Client) transcribeAudio(ctx context.Context, audio []byte, filename string, in Input) (*whisperResponse, time.Duration, error) {
 	body := &bytes.Buffer{}
 	mw := multipart.NewWriter(body)

internal/transcription/client_test.go

@@ -0,0 +1,28 @@
+package transcription
+
+import (
+	"math"
+	"testing"
+	"time"
+)
+
+func TestAdjustLeadSilence(t *testing.T) {
+	got := adjustLeadSilence([]Segment{
+		{Start: 0.2, End: 1.1, Text: "first"},
+		{Start: 1.4, End: 2.0, Text: "second"},
+	}, 800*time.Millisecond)
+
+	if got[0].Start != 0 {
+		t.Fatalf("first start = %v, want 0", got[0].Start)
+	}
+	if !near(got[0].End, 0.3) {
+		t.Fatalf("first end = %v, want 0.3", got[0].End)
+	}
+	if !near(got[1].Start, 0.6) {
+		t.Fatalf("second start = %v, want 0.6", got[1].Start)
+	}
+}
+
+func near(got, want float64) bool {
+	return math.Abs(got-want) < 0.000001
+}