Added SetFloat32 and SetTs to frame

Author: djthorpe

pkg/ffmpeg/frame.go

@@ -168,6 +168,29 @@ func (frame *Frame) Float32(plane int) []float32 {
 	return ctx.Float32(plane)
 }
 
+// Set plane data from float32 slice
+func (frame *Frame) SetFloat32(plane int, data []float32) error {
+	if frame.Type() != media.AUDIO {
+		return errors.New("frame is not an audio frame")
+	}
+
+	// If the number of samples is not the same, the re-allocate the frame
+	ctx := (*ff.AVFrame)(frame)
+	if len(data) != frame.NumSamples() {
+		ctx.SetNumSamples(len(data))
+		if err := ff.AVUtil_frame_get_buffer(ctx, false); err != nil {
+			ff.AVUtil_frame_unref(ctx)
+			return err
+		}
+	}
+
+	// Copy data
+	copy(ctx.Float32(plane), data)
+
+	// Return success
+	return nil
+}
+
 // Return plane data as a byte slice
 func (frame *Frame) Bytes(plane int) []byte {
 	return (*ff.AVFrame)(frame).Bytes(plane)
@@ -261,6 +284,17 @@ func (frame *Frame) Ts() float64 {
 	return ff.AVUtil_rational_q2d(tb) * float64(pts)
 }
 
+// Set timestamp in seconds
+func (frame *Frame) SetTs(secs float64) {
+	ctx := (*ff.AVFrame)(frame)
+	tb := ctx.TimeBase()
+	if secs == TS_UNDEFINED || tb.Num() == 0 || tb.Den() == 0 {
+		frame.SetPts(ff.AV_NOPTS_VALUE)
+		return
+	}
+	ctx.SetPts(int64(secs / ff.AVUtil_rational_q2d(tb)))
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
 

pkg/ffmpeg/frame_test.go

@@ -69,3 +69,44 @@ func Test_frame_004(t *testing.T) {
 	assert.Equal(copy.Width(), frame.Width())
 	assert.Equal(copy.Height(), frame.Height())
 }
+
+func Test_frame_005(t *testing.T) {
+	assert := assert.New(t)
+
+	frame, err := ffmpeg.NewFrame(ffmpeg.AudioPar("fltp", "stereo", 16000))
+	if !assert.NoError(err) {
+		t.FailNow()
+	}
+	defer frame.Close()
+
+	t.Log(frame)
+	data := make([]float32, frame.SampleRate())
+	frame.SetFloat32(0, data)
+	frame.SetFloat32(1, data)
+	t.Log(frame)
+
+	frame.SetFloat32(0, data[0:100])
+	frame.SetFloat32(1, data[0:100])
+	t.Log(frame)
+
+}
+
+func Test_frame_006(t *testing.T) {
+	assert := assert.New(t)
+
+	frame, err := ffmpeg.NewFrame(ffmpeg.AudioPar("fltp", "stereo", 16000))
+	if !assert.NoError(err) {
+		t.FailNow()
+	}
+	defer frame.Close()
+
+	for i := 0; i < 10; i++ {
+		frame.SetPts(int64(i))
+		assert.Equal(int64(i), frame.Pts())
+	}
+
+	for i := 0; i < 10; i++ {
+		frame.SetTs(float64(i))
+		assert.Equal(float64(i), frame.Ts())
+	}
+}