Skip to content

Commit bea7ae3

Browse files
kevmo314kevmo314
authored
Add ivfwriter support for VP9
Adds the necessary wiring to get VP9 to work with `ivfwriter`. Update the README of save-to-disk to inform users it supports both VP8 and VP9. ivfwriter currently assumes 30 fps but it seems that the other codecs also assume 30 fps so that is not a net-new assumption.
1 parent 306dc37 commit bea7ae3

File tree

5 files changed

+122
-42
lines changed

5 files changed

+122
-42
lines changed

‎examples/play-from-disk-renegotiation/main.go

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ package main
1010
import (
1111
"encoding/json"
1212
"fmt"
13-
"math/rand"
1413
"net/http"
1514
"os"
1615
"time"
@@ -116,8 +115,6 @@ func removeVideo(res http.ResponseWriter, req *http.Request) {
116115
}
117116

118117
func main() {
119-
rand.Seed(time.Now().UTC().UnixNano())
120-
121118
var err error
122119
if peerConnection, err = webrtc.NewPeerConnection(webrtc.Configuration{}); err != nil {
123120
panic(err)

‎examples/save-to-disk/README.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
# save-to-disk
22
save-to-disk is a simple application that shows how to record your webcam/microphone using Pion WebRTC and save VP8/Opus to disk.
33

4+
If you wish to save VP9 instead of VP8 you can just replace all occurences of VP8 with VP9 in [main.go](https://github.com/pion/example-webrtc-applications/tree/master/save-to-disk/main.go).
5+
46
If you wish to save VP8/Opus inside the same file see [save-to-webm](https://github.com/pion/example-webrtc-applications/tree/master/save-to-webm)
57

68
If you wish to save AV1 instead see [save-to-disk-av1](https://github.com/pion/webrtc/tree/master/examples/save-to-disk-av1)

‎examples/save-to-disk/main.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ func main() {
124124
if err != nil {
125125
panic(err)
126126
}
127-
ivfFile, err := ivfwriter.New("output.ivf")
127+
ivfFile, err := ivfwriter.New("output.ivf", ivfwriter.WithCodec("video/VP8"))
128128
if err != nil {
129129
panic(err)
130130
}

‎pkg/media/ivfwriter/ivfwriter.go

Lines changed: 84 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -16,40 +16,48 @@ import (
1616
)
1717

1818
var (
19-
errFileNotOpened = errors.New("file not opened")
20-
errInvalidNilPacket = errors.New("invalid nil packet")
21-
errCodecAlreadySet = errors.New("codec is already set")
22-
errNoSuchCodec = errors.New("no codec for this MimeType")
19+
errFileNotOpened = errors.New("file not opened")
20+
errInvalidNilPacket = errors.New("invalid nil packet")
21+
errCodecUnset = errors.New("codec is unset")
22+
errCodecAlreadySet = errors.New("codec is already set")
23+
errNoSuchCodec = errors.New("no codec for this MimeType")
24+
errInvalidMediaTimebase = errors.New("invalid media timebase")
2325
)
2426

25-
const (
26-
mimeTypeVP8 = "video/VP8"
27-
mimeTypeAV1 = "video/AV1"
27+
type (
28+
codec int
2829

29-
ivfFileHeaderSignature = "DKIF"
30-
)
30+
// IVFWriter is used to take RTP packets and write them to an IVF on disk.
31+
IVFWriter struct {
32+
ioWriter io.Writer
33+
count uint64
34+
seenKeyFrame bool
3135

32-
var errInvalidMediaTimebase = errors.New("invalid media timebase")
36+
codec codec
3337

34-
// IVFWriter is used to take RTP packets and write them to an IVF on disk.
35-
type IVFWriter struct {
36-
ioWriter io.Writer
37-
count uint64
38-
seenKeyFrame bool
38+
timebaseDenominator uint32
39+
timebaseNumerator uint32
40+
firstFrameTimestamp uint32
41+
clockRate uint64
3942

40-
isVP8, isAV1 bool
43+
// VP8, VP9
44+
currentFrame []byte
4145

42-
timebaseDenominator uint32
43-
timebaseNumerator uint32
44-
firstFrameTimestamp uint32
45-
clockRate uint64
46+
// AV1
47+
av1Frame frame.AV1
48+
}
49+
)
4650

47-
// VP8
48-
currentFrame []byte
51+
const (
52+
codecUnset codec = iota
53+
codecVP8
54+
codecVP9
55+
codecAV1
4956

50-
// AV1
51-
av1Frame frame.AV1
52-
}
57+
mimeTypeVP8 = "video/VP8"
58+
mimeTypeVP9 = "video/VP9"
59+
mimeTypeAV1 = "video/AV1"
60+
)
5361

5462
// New builds a new IVF writer.
5563
func New(fileName string, opts ...Option) (*IVFWriter, error) {
@@ -86,8 +94,8 @@ func NewWith(out io.Writer, opts ...Option) (*IVFWriter, error) {
8694
}
8795
}
8896

89-
if !writer.isAV1 && !writer.isVP8 {
90-
writer.isVP8 = true
97+
if writer.codec == codecUnset {
98+
writer.codec = codecVP8
9199
}
92100

93101
if err := writer.writeHeader(); err != nil {
@@ -103,15 +111,20 @@ func NewWith(out io.Writer, opts ...Option) (*IVFWriter, error) {
103111

104112
func (i *IVFWriter) writeHeader() error {
105113
header := make([]byte, 32)
106-
copy(header[0:], ivfFileHeaderSignature) // DKIF
114+
copy(header[0:], "DKIF") // DKIF
107115
binary.LittleEndian.PutUint16(header[4:], 0) // Version
108116
binary.LittleEndian.PutUint16(header[6:], 32) // Header size
109117

110118
// FOURCC
111-
if i.isVP8 {
119+
switch i.codec {
120+
case codecVP8:
112121
copy(header[8:], "VP80")
113-
} else if i.isAV1 {
122+
case codecVP9:
123+
copy(header[8:], "VP90")
124+
case codecAV1:
114125
copy(header[8:], "AV01")
126+
default:
127+
return errCodecUnset
115128
}
116129

117130
binary.LittleEndian.PutUint16(header[12:], 640) // Width in pixels
@@ -146,19 +159,20 @@ func (i *IVFWriter) writeFrame(frame []byte, timestamp uint64) error {
146159
}
147160

148161
// WriteRTP adds a new packet and writes the appropriate headers for it.
149-
func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error { //nolint:cyclop
162+
func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error { //nolint:cyclop, gocognit
150163
if i.ioWriter == nil {
151164
return errFileNotOpened
152165
} else if len(packet.Payload) == 0 {
153166
return nil
154167
}
155168

156169
if i.count == 0 {
157-
i.firstFrameTimestamp = packet.Header.Timestamp
170+
i.firstFrameTimestamp = packet.Timestamp
158171
}
159-
relativeTstampMs := 1000 * uint64(packet.Header.Timestamp-i.firstFrameTimestamp) / i.clockRate
172+
relativeTstampMs := 1000 * uint64(packet.Timestamp-i.firstFrameTimestamp) / i.clockRate
160173

161-
if i.isVP8 { //nolint:nestif
174+
switch i.codec {
175+
case codecVP8:
162176
vp8Packet := codecs.VP8Packet{}
163177
if _, err := vp8Packet.Unmarshal(packet.Payload); err != nil {
164178
return err
@@ -185,7 +199,35 @@ func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error { //nolint:cyclop
185199
return err
186200
}
187201
i.currentFrame = nil
188-
} else if i.isAV1 {
202+
case codecVP9:
203+
vp9Packet := codecs.VP9Packet{}
204+
if _, err := vp9Packet.Unmarshal(packet.Payload); err != nil {
205+
return err
206+
}
207+
208+
switch {
209+
case !i.seenKeyFrame && vp9Packet.P:
210+
return nil
211+
case i.currentFrame == nil && !vp9Packet.B:
212+
return nil
213+
}
214+
215+
i.seenKeyFrame = true
216+
i.currentFrame = append(i.currentFrame, vp9Packet.Payload[0:]...)
217+
218+
if !packet.Marker {
219+
return nil
220+
} else if len(i.currentFrame) == 0 {
221+
return nil
222+
}
223+
224+
// the timestamp must be sequential. webrtc mandates a clock rate of 90000
225+
// and we've assumed 30fps in the header.
226+
if err := i.writeFrame(i.currentFrame, uint64(packet.Timestamp)/3000); err != nil {
227+
return err
228+
}
229+
i.currentFrame = nil
230+
case codecAV1:
189231
av1Packet := &codecs.AV1Packet{}
190232
if _, err := av1Packet.Unmarshal(packet.Payload); err != nil {
191233
return err
@@ -201,6 +243,8 @@ func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error { //nolint:cyclop
201243
return err
202244
}
203245
}
246+
default:
247+
return errCodecUnset
204248
}
205249

206250
return nil
@@ -243,15 +287,17 @@ type Option func(i *IVFWriter) error
243287
// WithCodec configures if IVFWriter is writing AV1 or VP8 packets to disk.
244288
func WithCodec(mimeType string) Option {
245289
return func(i *IVFWriter) error {
246-
if i.isVP8 || i.isAV1 {
290+
if i.codec != codecUnset {
247291
return errCodecAlreadySet
248292
}
249293

250294
switch mimeType {
251295
case mimeTypeVP8:
252-
i.isVP8 = true
296+
i.codec = codecVP8
297+
case mimeTypeVP9:
298+
i.codec = codecVP9
253299
case mimeTypeAV1:
254-
i.isAV1 = true
300+
i.codec = codecAV1
255301
default:
256302
return errNoSuchCodec
257303
}

‎pkg/media/ivfwriter/ivfwriter_test.go

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -302,3 +302,38 @@ func TestIVFWriter_AV1(t *testing.T) {
302302
assert.NoError(t, writer.Close())
303303
})
304304
}
305+
306+
func TestIVFWriter_VP9(t *testing.T) {
307+
buffer := &bytes.Buffer{}
308+
writer, err := NewWith(buffer, WithCodec(mimeTypeVP9))
309+
assert.NoError(t, err)
310+
311+
// No keyframe yet, ignore non-keyframe packets (P)
312+
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0xD0, 0x02, 0xAA}}))
313+
assert.Equal(t, buffer.Bytes(), []byte{
314+
0x44, 0x4b, 0x49, 0x46, 0x00, 0x00, 0x20, 0x00, 0x56, 0x50, 0x39, 0x30, 0x80, 0x02, 0xe0, 0x01,
315+
0x1e, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x84, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
316+
})
317+
318+
// No current frame, ignore packets that don't start a frame (B)
319+
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x00, 0xAA}}))
320+
assert.Equal(t, buffer.Bytes(), []byte{
321+
0x44, 0x4b, 0x49, 0x46, 0x00, 0x00, 0x20, 0x00, 0x56, 0x50, 0x39, 0x30, 0x80, 0x02, 0xe0, 0x01,
322+
0x1e, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x84, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
323+
})
324+
325+
// B packet, no marker bit
326+
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x08, 0xAA}}))
327+
assert.Equal(t, buffer.Bytes(), []byte{
328+
0x44, 0x4b, 0x49, 0x46, 0x00, 0x00, 0x20, 0x00, 0x56, 0x50, 0x39, 0x30, 0x80, 0x02, 0xe0, 0x01,
329+
0x1e, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x84, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
330+
})
331+
332+
// B packet, Marker Bit
333+
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Header: rtp.Header{Marker: true}, Payload: []byte{0x08, 0xAB}}))
334+
assert.Equal(t, buffer.Bytes(), []byte{
335+
0x44, 0x4b, 0x49, 0x46, 0x00, 0x00, 0x20, 0x00, 0x56, 0x50, 0x39, 0x30, 0x80, 0x02, 0xe0, 0x01,
336+
0x1e, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x84, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
337+
0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xaa, 0xab,
338+
})
339+
}

0 commit comments

Comments
 (0)