Skip to content

Commit fd7fe3f

Browse files
JoeTurkiJoeTurki
committed
Fix AV1 IVF writer using RTP depacketizer
1. Replaced the OBU reader with the AV1 depacketizer to actually handle RTP input. 2. Now explicitly constructing OBUTemporalDelimiter to ensure proper separation of frames.
1 parent 38fdb70 commit fd7fe3f

File tree

2 files changed

+93
-25
lines changed

2 files changed

+93
-25
lines changed

‎pkg/media/ivfwriter/ivfwriter.go‎

Lines changed: 30 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import (
1212

1313
"github.com/pion/rtp"
1414
"github.com/pion/rtp/codecs"
15-
"github.com/pion/rtp/codecs/av1/frame"
15+
"github.com/pion/rtp/codecs/av1/obu"
1616
)
1717

1818
var (
@@ -44,7 +44,7 @@ type (
4444
currentFrame []byte
4545

4646
// AV1
47-
av1Frame frame.AV1
47+
av1Depacketizer *codecs.AV1Depacketizer
4848
}
4949
)
5050

@@ -97,7 +97,6 @@ func NewWith(out io.Writer, opts ...Option) (*IVFWriter, error) {
9797
if writer.codec == codecUnset {
9898
writer.codec = codecVP8
9999
}
100-
101100
if err := writer.writeHeader(); err != nil {
102101
return nil, err
103102
}
@@ -159,7 +158,7 @@ func (i *IVFWriter) writeFrame(frame []byte, timestamp uint64) error {
159158
}
160159

161160
// WriteRTP adds a new packet and writes the appropriate headers for it.
162-
func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error { //nolint:cyclop, gocognit
161+
func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error { //nolint:cyclop,gocyclo,gocognit
163162
if i.ioWriter == nil {
164163
return errFileNotOpened
165164
} else if len(packet.Payload) == 0 {
@@ -228,21 +227,41 @@ func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error { //nolint:cyclop, gocogn
228227
}
229228
i.currentFrame = nil
230229
case codecAV1:
231-
av1Packet := &codecs.AV1Packet{}
232-
if _, err := av1Packet.Unmarshal(packet.Payload); err != nil {
233-
return err
230+
if i.av1Depacketizer == nil {
231+
i.av1Depacketizer = &codecs.AV1Depacketizer{}
234232
}
235233

236-
obus, err := i.av1Frame.ReadFrames(av1Packet)
234+
payload, err := i.av1Depacketizer.Unmarshal(packet.Payload)
237235
if err != nil {
238236
return err
239237
}
240238

241-
for j := range obus {
242-
if err := i.writeFrame(obus[j], relativeTstampMs); err != nil {
243-
return err
239+
if !i.seenKeyFrame {
240+
isKeyFrame := i.av1Depacketizer.N || (len(payload) > 0 && obu.Type((payload[0]&0x78)>>3) == obu.OBUSequenceHeader)
241+
if !isKeyFrame {
242+
return nil
244243
}
244+
245+
i.seenKeyFrame = true
246+
}
247+
248+
i.currentFrame = append(i.currentFrame, payload...)
249+
250+
if !packet.Marker {
251+
return nil
252+
}
253+
254+
delimiter := obu.Header{
255+
Type: obu.OBUTemporalDelimiter,
256+
HasSizeField: true,
245257
}
258+
frame := append(delimiter.Marshal(), 0)
259+
frame = append(frame, i.currentFrame...)
260+
261+
if err := i.writeFrame(frame, relativeTstampMs); err != nil {
262+
return err
263+
}
264+
i.currentFrame = nil
246265
default:
247266
return errCodecUnset
248267
}

‎pkg/media/ivfwriter/ivfwriter_test.go‎

Lines changed: 63 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -263,15 +263,22 @@ func TestIVFWriter_AV1(t *testing.T) {
263263
writer, err := NewWith(buffer, WithCodec(mimeTypeAV1))
264264
assert.NoError(t, err)
265265

266-
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x00, 0x01, 0xFF}}))
266+
assert.NoError(
267+
t,
268+
writer.WriteRTP(
269+
&rtp.Packet{
270+
Header: rtp.Header{Marker: true},
271+
// N = 1, Length = 1, OBU_TYPE = 4
272+
Payload: []byte{0x08, 0x01, 0x20},
273+
}),
274+
)
275+
267276
assert.NoError(t, writer.Close())
268277
assert.Equal(t, buffer.Bytes(), []byte{
269-
0x44, 0x4b, 0x49, 0x46, 0x0, 0x0, 0x20,
270-
0x0, 0x41, 0x56, 0x30, 0x31, 0x80, 0x2,
271-
0xe0, 0x1, 0x1e, 0x0, 0x0, 0x0, 0x1, 0x0,
272-
0x0, 0x0, 0x84, 0x3, 0x0, 0x0, 0x0, 0x0,
273-
0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
274-
0x0, 0x0, 0x0, 0x0, 0x0, 0xff,
278+
0x44, 0x4b, 0x49, 0x46, 0x0, 0x0, 0x20, 0x0, 0x41, 0x56, 0x30, 0x31,
279+
0x80, 0x2, 0xe0, 0x1, 0x1e, 0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x84,
280+
0x3, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x4, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
281+
0x0, 0x0, 0x0, 0x0, 0x0, 0x12, 0x0, 0x22, 0x0,
275282
})
276283
})
277284

@@ -281,8 +288,13 @@ func TestIVFWriter_AV1(t *testing.T) {
281288
writer, err := NewWith(buffer, WithCodec(mimeTypeAV1))
282289
assert.NoError(t, err)
283290

284-
for _, p := range [][]byte{{0x40, 0x02, 0x00, 0x01}, {0xc0, 0x02, 0x02, 0x03}, {0xc0, 0x02, 0x04, 0x04}} {
285-
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: p}))
291+
for _, p := range [][]byte{
292+
{0x48, 0x02, 0x00, 0x01}, // Y=true
293+
{0xc0, 0x02, 0x02, 0x03}, // Z=true, Y=true
294+
{0xc0, 0x02, 0x04, 0x04}, // Z=true, Y=true
295+
{0x80, 0x01, 0x05}, // Z=true, Y=false (But we still don't set Marker to true)
296+
} {
297+
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: p, Header: rtp.Header{Marker: false}}))
286298
assert.Equal(t, buffer.Bytes(), []byte{
287299
0x44, 0x4b, 0x49, 0x46, 0x0,
288300
0x0, 0x20, 0x0, 0x41, 0x56, 0x30,
@@ -292,15 +304,52 @@ func TestIVFWriter_AV1(t *testing.T) {
292304
0x0, 0x0,
293305
})
294306
}
295-
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x80, 0x01, 0x05}}))
307+
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x08, 0x01, 0x20}, Header: rtp.Header{Marker: true}}))
296308
assert.Equal(t, buffer.Bytes(), []byte{
297-
0x44, 0x4b, 0x49, 0x46, 0x0, 0x0, 0x20, 0x0, 0x41, 0x56, 0x30, 0x31, 0x80,
298-
0x2, 0xe0, 0x1, 0x1e, 0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x84, 0x3, 0x0, 0x0,
299-
0x0, 0x0, 0x0, 0x0, 0x7, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
300-
0x0, 0x1, 0x2, 0x3, 0x4, 0x4, 0x5,
309+
0x44, 0x4b, 0x49, 0x46, 0x0, 0x0, 0x20, 0x0, 0x41, 0x56, 0x30, 0x31, 0x80, 0x2, 0xe0, 0x1, 0x1e,
310+
0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x84, 0x3, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0xc, 0x0, 0x0, 0x0,
311+
0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x12, 0x0, 0x2, 0x6, 0x1, 0x2, 0x3, 0x4, 0x4, 0x5, 0x22, 0x0,
301312
})
302313
assert.NoError(t, writer.Close())
303314
})
315+
316+
t.Run("Invalid OBU", func(t *testing.T) {
317+
buffer := &bytes.Buffer{}
318+
319+
writer, err := NewWith(buffer, WithCodec(mimeTypeAV1))
320+
assert.NoError(t, err)
321+
322+
assert.Error(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x08, 0x02, 0xff}}))
323+
assert.Error(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x08, 0x01, 0xff}}))
324+
})
325+
326+
t.Run("Skips middle sequence start", func(t *testing.T) {
327+
buffer := &bytes.Buffer{}
328+
329+
writer, err := NewWith(buffer, WithCodec(mimeTypeAV1))
330+
assert.NoError(t, err)
331+
332+
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Header: rtp.Header{Marker: true}, Payload: []byte{0x00, 0x01, 0x20}}))
333+
334+
assert.NoError(
335+
t,
336+
writer.WriteRTP(
337+
&rtp.Packet{
338+
Header: rtp.Header{Marker: true},
339+
// N = 1, Length = 1, OBU_TYPE = 4
340+
Payload: []byte{0x08, 0x01, 0x20},
341+
},
342+
),
343+
)
344+
345+
assert.NoError(t, writer.Close())
346+
assert.Equal(t, buffer.Bytes(), []byte{
347+
0x44, 0x4b, 0x49, 0x46, 0x0, 0x0, 0x20, 0x0, 0x41, 0x56, 0x30, 0x31,
348+
0x80, 0x2, 0xe0, 0x1, 0x1e, 0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x84,
349+
0x3, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x4, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
350+
0x0, 0x0, 0x0, 0x0, 0x0, 0x12, 0x0, 0x22, 0x0,
351+
})
352+
})
304353
}
305354

306355
func TestIVFWriter_VP9(t *testing.T) {

0 commit comments

Comments
 (0)