diff --git a/pkg/ffmpeg/frame.go b/pkg/ffmpeg/frame.go index 365b67b..56d385e 100644 --- a/pkg/ffmpeg/frame.go +++ b/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 diff --git a/pkg/ffmpeg/frame_test.go b/pkg/ffmpeg/frame_test.go index 894fae3..07df710 100644 --- a/pkg/ffmpeg/frame_test.go +++ b/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()) + } +}