Skip to content

Commit 9c2694f

Browse files
committed
Added encoding
1 parent 61ee559 commit 9c2694f

File tree

8 files changed

+166
-68
lines changed

8 files changed

+166
-68
lines changed

pkg/ffmpeg/encoder.go

Lines changed: 27 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package ffmpeg
33
import (
44
"encoding/json"
55
"errors"
6+
"fmt"
67
"io"
78
"syscall"
89

@@ -23,7 +24,9 @@ type Encoder struct {
2324

2425
// We are flushing the encoder
2526
eof bool
26-
//next_pts int64
27+
28+
// The next presentation timestamp
29+
next_pts int64
2730
}
2831

2932
// EncoderFrameFn is a function which is called to receive a frame to encode. It should
@@ -85,12 +88,6 @@ func NewEncoder(ctx *ff.AVFormatContext, stream int, par *Par) (*Encoder, error)
8588
encoder.stream = streamctx
8689
}
8790

88-
// Copy parameters to stream
89-
if err := ff.AVCodec_parameters_from_context(encoder.stream.CodecPar(), encoder.ctx); err != nil {
90-
ff.AVCodec_free_context(encoder.ctx)
91-
return nil, err
92-
}
93-
9491
// Some formats want stream headers to be separate.
9592
if ctx.Output().Flags().Is(ff.AVFMT_GLOBALHEADER) {
9693
encoder.ctx.SetFlags(encoder.ctx.Flags() | ff.AV_CODEC_FLAG_GLOBAL_HEADER)
@@ -102,6 +99,12 @@ func NewEncoder(ctx *ff.AVFormatContext, stream int, par *Par) (*Encoder, error)
10299
return nil, ErrInternalAppError.Withf("codec_open: %v", err)
103100
}
104101

102+
// Copy parameters to stream
103+
if err := ff.AVCodec_parameters_from_context(encoder.stream.CodecPar(), encoder.ctx); err != nil {
104+
ff.AVCodec_free_context(encoder.ctx)
105+
return nil, err
106+
}
107+
105108
// Create a packet
106109
packet := ff.AVCodec_packet_alloc()
107110
if packet == nil {
@@ -176,6 +179,22 @@ func (e *Encoder) Par() *Par {
176179
}
177180
}
178181

182+
// Return the codec type
183+
func (e *Encoder) nextPts(frame *ff.AVFrame) int64 {
184+
next_pts := int64(0)
185+
switch e.ctx.Codec().Type() {
186+
case ff.AVMEDIA_TYPE_AUDIO:
187+
next_pts = ff.AVUtil_rational_rescale_q(int64(frame.NumSamples()), ff.AVUtil_rational(1, frame.SampleRate()), e.stream.TimeBase())
188+
case ff.AVMEDIA_TYPE_VIDEO:
189+
next_pts = ff.AVUtil_rational_rescale_q(1, ff.AVUtil_rational_invert(e.ctx.Framerate()), e.stream.TimeBase())
190+
default:
191+
// Dunno what to do with subtitle and data streams yet
192+
fmt.Println("TODO: next_pts for subtitle and data streams")
193+
return 0
194+
}
195+
return next_pts
196+
}
197+
179198
//////////////////////////////////////////////////////////////////////////////
180199
// PRIVATE METHODS
181200

@@ -201,6 +220,7 @@ func (e *Encoder) encode(frame *ff.AVFrame, fn EncoderPacketFn) error {
201220
// rescale output packet timestamp values from codec to stream timebase
202221
ff.AVCodec_packet_rescale_ts(e.packet, e.ctx.TimeBase(), e.stream.TimeBase())
203222
e.packet.SetStreamIndex(e.stream.Index())
223+
e.packet.SetTimeBase(e.stream.TimeBase())
204224

205225
// Pass back to the caller
206226
if err := fn(e.packet, &timebase); errors.Is(err, io.EOF) {

pkg/ffmpeg/rescaler_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ func Test_rescaler_001(t *testing.T) {
1717
assert := assert.New(t)
1818

1919
// Create an image generator
20-
image, err := generator.NewYUV420P(25, ffmpeg.VideoPar("yuv420p", "1280x720", 25))
20+
image, err := generator.NewYUV420P(ffmpeg.VideoPar("yuv420p", "1280x720", 25))
2121
if !assert.NoError(err) {
2222
t.FailNow()
2323
}
@@ -53,7 +53,7 @@ func Test_rescaler_002(t *testing.T) {
5353
assert := assert.New(t)
5454

5555
// Create an image generator
56-
image, err := generator.NewYUV420P(25, ffmpeg.VideoPar("yuva420p", "1280x720", 25))
56+
image, err := generator.NewYUV420P(ffmpeg.VideoPar("yuva420p", "1280x720", 25))
5757
if !assert.NoError(err) {
5858
t.FailNow()
5959
}

pkg/ffmpeg/writer.go

Lines changed: 36 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,7 @@ func (w *Writer) Encode(in EncoderFrameFn, out EncoderPacketFn) error {
230230

231231
// Initialize the encoder
232232
encoder.eof = false
233+
encoder.next_pts = 0
233234
}
234235

235236
// Continue until all encoders have returned io.EOF and have been flushed
@@ -239,31 +240,42 @@ func (w *Writer) Encode(in EncoderFrameFn, out EncoderPacketFn) error {
239240
break
240241
}
241242

242-
// TODO: We get the encoder with the lowest timestamp
243+
// Find encoder with the lowest timestamp, based on next_pts and timebase
244+
next_stream := -1
245+
var next_encoder *Encoder
243246
for stream, encoder := range encoders {
244-
var frame *ff.AVFrame
245-
var err error
246-
247-
// Receive a frame if not EOF
248-
if !encoder.eof {
249-
frame, err = in(stream)
250-
if errors.Is(err, io.EOF) {
251-
encoder.eof = true
252-
} else if err != nil {
253-
return fmt.Errorf("stream %v: %w", stream, err)
254-
}
247+
if next_encoder == nil || compareNextPts(encoder, next_encoder) < 0 {
248+
next_encoder = encoder
249+
next_stream = stream
255250
}
251+
}
256252

257-
// Send a frame for encoding
258-
if err := encoder.Encode(frame, out); err != nil {
259-
return fmt.Errorf("stream %v: %w", stream, err)
260-
}
253+
var frame *ff.AVFrame
254+
var err error
261255

262-
// If eof then delete the encoder
263-
if encoder.eof {
264-
delete(encoders, stream)
256+
// Receive a frame if not EOF
257+
if !next_encoder.eof {
258+
frame, err = in(next_stream)
259+
if errors.Is(err, io.EOF) {
260+
next_encoder.eof = true
261+
} else if err != nil {
262+
return fmt.Errorf("stream %v: %w", next_stream, err)
265263
}
266264
}
265+
266+
// Send a frame for encoding
267+
if err := next_encoder.Encode(frame, out); err != nil {
268+
return fmt.Errorf("stream %v: %w", next_stream, err)
269+
}
270+
271+
// If eof then delete the encoder
272+
if next_encoder.eof {
273+
delete(encoders, next_stream)
274+
continue
275+
}
276+
277+
// Calculate the next PTS
278+
next_encoder.next_pts = next_encoder.next_pts + next_encoder.nextPts(frame)
267279
}
268280

269281
// Return success
@@ -276,6 +288,11 @@ func (w *Writer) Write(packet *ff.AVPacket) error {
276288
return ff.AVCodec_interleaved_write_frame(w.output, packet)
277289
}
278290

291+
// Returns -1 if a is before v
292+
func compareNextPts(a, b *Encoder) int {
293+
return ff.AVUtil_compare_ts(a.next_pts, a.stream.TimeBase(), b.next_pts, b.stream.TimeBase())
294+
}
295+
279296
////////////////////////////////////////////////////////////////////////////////
280297
// PRIVATE METHODS - Writer
281298

pkg/ffmpeg/writer_test.go

Lines changed: 61 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"testing"
77
"time"
88

9+
media "github.com/mutablelogic/go-media"
910
ffmpeg "github.com/mutablelogic/go-media/pkg/ffmpeg"
1011
generator "github.com/mutablelogic/go-media/pkg/generator"
1112
ff "github.com/mutablelogic/go-media/sys/ffmpeg61"
@@ -118,15 +119,15 @@ func Test_writer_003(t *testing.T) {
118119
// Create a writer with an audio stream
119120
writer, err := ffmpeg.NewWriter(w,
120121
ffmpeg.OptMetadata(ffmpeg.NewMetadata("title", t.Name())),
121-
ffmpeg.OptStream(1, ffmpeg.VideoPar("yuv420p", "1280x720", 25)),
122+
ffmpeg.OptStream(1, ffmpeg.VideoPar("yuv420p", "640x480", 30)),
122123
)
123124
if !assert.NoError(err) {
124125
t.FailNow()
125126
}
126127
defer writer.Close()
127128

128129
// Make an video generator
129-
video, err := generator.NewYUV420P(25, writer.Stream(1).Par())
130+
video, err := generator.NewYUV420P(writer.Stream(1).Par())
130131
if !assert.NoError(err) {
131132
t.FailNow()
132133
}
@@ -139,14 +140,70 @@ func Test_writer_003(t *testing.T) {
139140
if frame.Time() >= duration {
140141
return nil, io.EOF
141142
} else {
142-
t.Log("Frame", frame.Time().Truncate(time.Millisecond))
143+
t.Log("Frame", stream, "=>", frame.Time().Truncate(time.Millisecond))
143144
return frame.(*ffmpeg.Frame).AVFrame(), nil
144145
}
145146
}, func(packet *ff.AVPacket, timebase *ff.AVRational) error {
146147
if packet != nil {
147-
t.Log("Packet", packet)
148+
d := time.Duration(ff.AVUtil_rational_q2d(packet.TimeBase()) * float64(packet.Pts()) * float64(time.Second))
149+
t.Log("Packet", d.Truncate(time.Millisecond))
148150
}
149151
return writer.Write(packet)
150152
}))
151153
t.Log("Written to", w.Name())
152154
}
155+
156+
func Test_writer_004(t *testing.T) {
157+
assert := assert.New(t)
158+
159+
// Write to a file
160+
w, err := os.CreateTemp("", t.Name()+"_*.m4v")
161+
if !assert.NoError(err) {
162+
t.FailNow()
163+
}
164+
defer w.Close()
165+
166+
// Create a writer with an audio stream
167+
writer, err := ffmpeg.Create(w.Name(),
168+
ffmpeg.OptMetadata(ffmpeg.NewMetadata("title", t.Name())),
169+
ffmpeg.OptStream(1, ffmpeg.VideoPar("yuv420p", "640x480", 30)),
170+
ffmpeg.OptStream(2, ffmpeg.AudioPar("fltp", "mono", 22050)),
171+
)
172+
if !assert.NoError(err) {
173+
t.FailNow()
174+
}
175+
defer writer.Close()
176+
177+
// Make an video generator
178+
video, err := generator.NewYUV420P(writer.Stream(1).Par())
179+
if !assert.NoError(err) {
180+
t.FailNow()
181+
}
182+
defer video.Close()
183+
184+
// Make an audio generator
185+
audio, err := generator.NewSine(440, -5, writer.Stream(2).Par())
186+
if !assert.NoError(err) {
187+
t.FailNow()
188+
}
189+
190+
// Write 10 secs of frames
191+
duration := time.Minute * 10
192+
assert.NoError(writer.Encode(func(stream int) (*ff.AVFrame, error) {
193+
var frame media.Frame
194+
switch stream {
195+
case 1:
196+
frame = video.Frame()
197+
case 2:
198+
frame = audio.Frame()
199+
}
200+
if frame.Time() >= duration {
201+
t.Log("Frame time is EOF", frame.Time())
202+
return nil, io.EOF
203+
} else {
204+
t.Log("Frame", stream, "=>", frame.Time().Truncate(time.Millisecond))
205+
return frame.(*ffmpeg.Frame).AVFrame(), nil
206+
}
207+
}, nil))
208+
t.Log("Written to", w.Name())
209+
}

pkg/generator/yuv420p.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,19 +24,19 @@ var _ Generator = (*yuv420p)(nil)
2424

2525
// Create a new video generator which generates YUV420P frames
2626
// of the specified size and framerate (in frames per second)
27-
func NewYUV420P(framerate int, par *ffmpeg.Par) (*yuv420p, error) {
27+
func NewYUV420P(par *ffmpeg.Par) (*yuv420p, error) {
2828
yuv420p := new(yuv420p)
2929

30-
// Check parameters
31-
if framerate <= 0 {
32-
return nil, errors.New("invalid framerate")
33-
}
3430
// Check parameters
3531
if par.CodecType() != ff.AVMEDIA_TYPE_VIDEO {
3632
return nil, errors.New("invalid codec type")
3733
} else if par.PixelFormat() != ff.AV_PIX_FMT_YUV420P {
3834
return nil, errors.New("invalid pixel format, only yuv420p is supported")
3935
}
36+
framerate := ff.AVUtil_rational_q2d(par.Framerate())
37+
if framerate <= 0 {
38+
return nil, errors.New("invalid framerate")
39+
}
4040

4141
// Create a frame
4242
frame := ff.AVUtil_frame_alloc()
@@ -48,7 +48,7 @@ func NewYUV420P(framerate int, par *ffmpeg.Par) (*yuv420p, error) {
4848
frame.SetWidth(par.Width())
4949
frame.SetHeight(par.Height())
5050
frame.SetSampleAspectRatio(par.SampleAspectRatio())
51-
frame.SetTimeBase(ff.AVUtil_rational(1, framerate))
51+
frame.SetTimeBase(ff.AVUtil_rational_invert(par.Framerate()))
5252
frame.SetPts(ff.AV_NOPTS_VALUE)
5353

5454
// Allocate buffer

sys/ffmpeg61/avcodec_packet.go

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,14 +18,15 @@ import "C"
1818
// TYPES
1919

2020
type jsonAVPacket struct {
21-
Pts int64 `json:"pts,omitempty"`
22-
Dts int64 `json:"dts,omitempty"`
23-
Size int `json:"size,omitempty"`
24-
StreamIndex int `json:"stream_index"` // Stream index starts at 0
25-
Flags int `json:"flags,omitempty"`
26-
SideDataElems int `json:"side_data_elems,omitempty"`
27-
Duration int64 `json:"duration,omitempty"`
28-
Pos int64 `json:"pos,omitempty"`
21+
Pts int64 `json:"pts,omitempty"`
22+
Dts int64 `json:"dts,omitempty"`
23+
Size int `json:"size,omitempty"`
24+
StreamIndex int `json:"stream_index"` // Stream index starts at 0
25+
Flags int `json:"flags,omitempty"`
26+
SideDataElems int `json:"side_data_elems,omitempty"`
27+
Duration int64 `json:"duration,omitempty"`
28+
TimeBase AVRational `json:"time_base,omitempty"`
29+
Pos int64 `json:"pos,omitempty"`
2930
}
3031

3132
////////////////////////////////////////////////////////////////////////////////
@@ -40,6 +41,7 @@ func (ctx *AVPacket) MarshalJSON() ([]byte, error) {
4041
Flags: int(ctx.flags),
4142
SideDataElems: int(ctx.side_data_elems),
4243
Duration: int64(ctx.duration),
44+
TimeBase: AVRational(ctx.time_base),
4345
Pos: int64(ctx.pos),
4446
})
4547
}
@@ -114,6 +116,14 @@ func (ctx *AVPacket) SetStreamIndex(index int) {
114116
ctx.stream_index = C.int(index)
115117
}
116118

119+
func (ctx *AVPacket) TimeBase() AVRational {
120+
return AVRational(ctx.time_base)
121+
}
122+
123+
func (ctx *AVPacket) SetTimeBase(tb AVRational) {
124+
ctx.time_base = C.AVRational(tb)
125+
}
126+
117127
func (ctx *AVPacket) Pts() int64 {
118128
return int64(ctx.pts)
119129
}

sys/ffmpeg61/avutil_math.go

Lines changed: 0 additions & 22 deletions
This file was deleted.

0 commit comments

Comments
 (0)