Upgraded segmenter

Author: djthorpe

go.mod

@@ -8,7 +8,7 @@ require (
 	github.com/djthorpe/go-tablewriter v0.0.8
 	github.com/go-audio/wav v1.1.0
 	github.com/mutablelogic/go-client v1.0.9
-	github.com/mutablelogic/go-media v1.6.7
+	github.com/mutablelogic/go-media v1.6.8
 	github.com/mutablelogic/go-server v1.4.13
 	github.com/stretchr/testify v1.9.0
 )

go.sum

@@ -26,8 +26,8 @@ github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6T
 github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
 github.com/mutablelogic/go-client v1.0.9 h1:Eh4sjQOFDldP/L3IizqkcOD3WigZR+u1VaHTUM4ujYw=
 github.com/mutablelogic/go-client v1.0.9/go.mod h1:VLyB8j8IBJSK/FXvvqhmq93PRWDKkyLu8R7V2Vudb6A=
-github.com/mutablelogic/go-media v1.6.7 h1:0hCr89EVJg7xw8ChABb7Cscr0UiZ1+Tl9xDXong0lu0=
-github.com/mutablelogic/go-media v1.6.7/go.mod h1:vWKq6QKqUQ+sAwfbU/DgakJGIk2Uq7ozH0qSxhysCkM=
+github.com/mutablelogic/go-media v1.6.8 h1:3v4povSQlOnvg9mHx6Bp9NVdCCjrNdDCjMHBGFHnVE8=
+github.com/mutablelogic/go-media v1.6.8/go.mod h1:HulNT0yyH63a3FRlbuzNDakhOypYrmtFVkHEXZjDgAY=
 github.com/mutablelogic/go-server v1.4.13 h1:k5LJJ/pCvyiw34UX341vRhliBOS6i7V65U/UICcOJOA=
 github.com/mutablelogic/go-server v1.4.13/go.mod h1:9nenPAohKu8bFoRgwHJh+3s8h0kLFjUAb8KZvT1TQNU=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

pkg/whisper/segmenter/segmenter.go

@@ -3,26 +3,50 @@ package segmenter
 import (
 	"context"
 	"errors"
-	"fmt"
 	"io"
+	"time"
 
 	// Packages
 	media "github.com/mutablelogic/go-media"
 	ffmpeg "github.com/mutablelogic/go-media/pkg/ffmpeg"
+
+	// Namespace imports
+	. "github.com/djthorpe/go-errors"
 )
 
 type Segmenter struct {
-	reader *ffmpeg.Reader
+	ts          time.Duration
+	sample_rate int
+	n           int
+	buf         []float32
+	reader      *ffmpeg.Reader
 }
 
+// SegmentFunc is a callback function which is called when a segment is ready
+// to be processed. The first argument is the timestamp of the segment.
+type SegmentFunc func(time.Duration, []float32)
+
 //////////////////////////////////////////////////////////////////////////////
 // LIFECYCLE
 
 // Create a new segmenter for "NumSamples" with a reader r
 // If NumSamples is zero then no segmenting is performed
-func NewSegmenter(r io.Reader) (*Segmenter, error) {
+func NewSegmenter(r io.Reader, dur time.Duration, sample_rate int) (*Segmenter, error) {
 	segmenter := new(Segmenter)
 
+	// Check arguments
+	if dur < 0 || sample_rate <= 0 {
+		return nil, ErrBadParameter.With("invalid duration or sample rate arguments")
+	} else {
+		segmenter.sample_rate = sample_rate
+	}
+
+	// Sample buffer is duration * sample rate
+	if dur > 0 {
+		segmenter.n = int(dur.Seconds()) * sample_rate
+		segmenter.buf = make([]float32, 0, int(dur.Seconds())*sample_rate)
+	}
+
 	// Open the file
 	media, err := ffmpeg.NewReader(r)
 	if err != nil {
@@ -42,6 +66,7 @@ func (s *Segmenter) Close() error {
 		result = errors.Join(result, s.reader.Close())
 	}
 	s.reader = nil
+	s.buf = nil
 
 	// Return any errors
 	return result
@@ -53,19 +78,51 @@ func (s *Segmenter) Close() error {
 // TODO: segments are output through a callback, with the samples and a timestamp
 // TODO: we could do some basic silence and voice detection to segment to ensure
 // we don't overtax the CPU/GPU with silence and non-speech
-// TODO: We have hard-coded the sample format, sample rate and number of channels
-// here. We should make this configurable
-func (s *Segmenter) Decode(ctx context.Context) error {
+func (s *Segmenter) Decode(ctx context.Context, fn SegmentFunc) error {
+	// Check input parameters
+	if fn == nil {
+		return ErrBadParameter.With("SegmentFunc is nil")
+	}
+
+	// Map function chooses the best audio stream
 	mapFunc := func(stream int, params *ffmpeg.Par) (*ffmpeg.Par, error) {
 		if stream == s.reader.BestStream(media.AUDIO) {
-			return ffmpeg.NewAudioPar("flt", "mono", 16000)
+			return ffmpeg.NewAudioPar("flt", "mono", s.sample_rate)
 		}
 		// Ignore no-audio streams
 		return nil, nil
 	}
-	return s.reader.Decode(ctx, mapFunc, func(stream int, frame *ffmpeg.Frame) error {
-		// Append float32 samples to buffer
-		fmt.Println("TODO: Implement Decode", frame)
+
+	// Decode samples and segment
+	if err := s.reader.Decode(ctx, mapFunc, func(stream int, frame *ffmpeg.Frame) error {
+		// We get null frames sometimes, ignore them
+		if frame == nil {
+			return nil
+		}
+
+		// Append float32 samples from plane 0 to buffer
+		s.buf = append(s.buf, frame.Float32(0)...)
+
+		// n != 0 and len(buf) >= n we have a segment to process
+		if s.n != 0 && len(s.buf) >= s.n {
+			fn(s.ts, s.buf)
+			// Clear the buffer
+			s.buf = s.buf[:0]
+			// Increment the timestamp
+			s.ts += time.Duration(float64(s.n)/float64(s.sample_rate)) * time.Second
+		}
+
+		// Continue processing
 		return nil
-	})
+	}); err != nil {
+		return err
+	}
+
+	// Output any remaining samples
+	if len(s.buf) > 0 {
+		fn(s.ts, s.buf)
+	}
+
+	// Return success
+	return nil
 }

pkg/whisper/segmenter/segmenter_test.go

@@ -4,28 +4,49 @@ import (
 	"context"
 	"os"
 	"testing"
+	"time"
 
 	// Packages
 	segmenter "github.com/mutablelogic/go-whisper/pkg/whisper/segmenter"
 	assert "github.com/stretchr/testify/assert"
 )
 
-const SAMPLE_EN = "../../../samples/jfk.wav"
-const SAMPLE_FR = "../../../samples/OlivierL.wav"
-const SAMPLE_DE = "../../../samples/ge-podcast.wav"
+const SAMPLE = "../../../samples/OlivierL.wav"
 
 func Test_segmenter_001(t *testing.T) {
 	assert := assert.New(t)
 
-	f, err := os.Open(SAMPLE_EN)
+	f, err := os.Open(SAMPLE)
 	if !assert.NoError(err) {
 		t.SkipNow()
 	}
-	segmenter, err := segmenter.NewSegmenter(f)
+	segmenter, err := segmenter.NewSegmenter(f, time.Second, 16000)
 	if !assert.NoError(err) {
 		t.SkipNow()
 	}
 	defer segmenter.Close()
 
-	assert.NoError(segmenter.Decode(context.Background()))
+	assert.NoError(segmenter.Decode(context.Background(), func(ts time.Duration, buf []float32) {
+		t.Log(ts, len(buf))
+	}))
+}
+
+func Test_segmenter_002(t *testing.T) {
+	assert := assert.New(t)
+
+	f, err := os.Open(SAMPLE)
+	if !assert.NoError(err) {
+		t.SkipNow()
+	}
+
+	// No segmentation, just output the audio
+	segmenter, err := segmenter.NewSegmenter(f, 0, 16000)
+	if !assert.NoError(err) {
+		t.SkipNow()
+	}
+	defer segmenter.Close()
+
+	assert.NoError(segmenter.Decode(context.Background(), func(ts time.Duration, buf []float32) {
+		t.Log(ts, len(buf))
+	}))
 }