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.
This commit is contained in:
Kevin Wang
2025-02-12 16:06:41 -05:00
committed by GitHub
parent 306dc37769
commit bea7ae3745
5 changed files with 125 additions and 45 deletions

View File

@@ -10,7 +10,6 @@ package main
import (
"encoding/json"
"fmt"
"math/rand"
"net/http"
"os"
"time"
@@ -116,8 +115,6 @@ func removeVideo(res http.ResponseWriter, req *http.Request) {
}
func main() {
rand.Seed(time.Now().UTC().UnixNano())
var err error
if peerConnection, err = webrtc.NewPeerConnection(webrtc.Configuration{}); err != nil {
panic(err)

View File

@@ -1,6 +1,8 @@
# save-to-disk
save-to-disk is a simple application that shows how to record your webcam/microphone using Pion WebRTC and save VP8/Opus to disk.
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).
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)
If you wish to save AV1 instead see [save-to-disk-av1](https://github.com/pion/webrtc/tree/master/examples/save-to-disk-av1)

View File

@@ -124,7 +124,7 @@ func main() {
if err != nil {
panic(err)
}
ivfFile, err := ivfwriter.New("output.ivf")
ivfFile, err := ivfwriter.New("output.ivf", ivfwriter.WithCodec("video/VP8"))
if err != nil {
panic(err)
}

View File

@@ -16,41 +16,49 @@ import (
)
var (
errFileNotOpened = errors.New("file not opened")
errInvalidNilPacket = errors.New("invalid nil packet")
errCodecAlreadySet = errors.New("codec is already set")
errNoSuchCodec = errors.New("no codec for this MimeType")
errFileNotOpened = errors.New("file not opened")
errInvalidNilPacket = errors.New("invalid nil packet")
errCodecUnset = errors.New("codec is unset")
errCodecAlreadySet = errors.New("codec is already set")
errNoSuchCodec = errors.New("no codec for this MimeType")
errInvalidMediaTimebase = errors.New("invalid media timebase")
)
type (
codec int
// IVFWriter is used to take RTP packets and write them to an IVF on disk.
IVFWriter struct {
ioWriter io.Writer
count uint64
seenKeyFrame bool
codec codec
timebaseDenominator uint32
timebaseNumerator uint32
firstFrameTimestamp uint32
clockRate uint64
// VP8, VP9
currentFrame []byte
// AV1
av1Frame frame.AV1
}
)
const (
codecUnset codec = iota
codecVP8
codecVP9
codecAV1
mimeTypeVP8 = "video/VP8"
mimeTypeVP9 = "video/VP9"
mimeTypeAV1 = "video/AV1"
ivfFileHeaderSignature = "DKIF"
)
var errInvalidMediaTimebase = errors.New("invalid media timebase")
// IVFWriter is used to take RTP packets and write them to an IVF on disk.
type IVFWriter struct {
ioWriter io.Writer
count uint64
seenKeyFrame bool
isVP8, isAV1 bool
timebaseDenominator uint32
timebaseNumerator uint32
firstFrameTimestamp uint32
clockRate uint64
// VP8
currentFrame []byte
// AV1
av1Frame frame.AV1
}
// New builds a new IVF writer.
func New(fileName string, opts ...Option) (*IVFWriter, error) {
file, err := os.Create(fileName) //nolint:gosec
@@ -86,8 +94,8 @@ func NewWith(out io.Writer, opts ...Option) (*IVFWriter, error) {
}
}
if !writer.isAV1 && !writer.isVP8 {
writer.isVP8 = true
if writer.codec == codecUnset {
writer.codec = codecVP8
}
if err := writer.writeHeader(); err != nil {
@@ -103,15 +111,20 @@ func NewWith(out io.Writer, opts ...Option) (*IVFWriter, error) {
func (i *IVFWriter) writeHeader() error {
header := make([]byte, 32)
copy(header[0:], ivfFileHeaderSignature) // DKIF
copy(header[0:], "DKIF") // DKIF
binary.LittleEndian.PutUint16(header[4:], 0) // Version
binary.LittleEndian.PutUint16(header[6:], 32) // Header size
// FOURCC
if i.isVP8 {
switch i.codec {
case codecVP8:
copy(header[8:], "VP80")
} else if i.isAV1 {
case codecVP9:
copy(header[8:], "VP90")
case codecAV1:
copy(header[8:], "AV01")
default:
return errCodecUnset
}
binary.LittleEndian.PutUint16(header[12:], 640) // Width in pixels
@@ -146,7 +159,7 @@ func (i *IVFWriter) writeFrame(frame []byte, timestamp uint64) error {
}
// WriteRTP adds a new packet and writes the appropriate headers for it.
func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error { //nolint:cyclop
func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error { //nolint:cyclop, gocognit
if i.ioWriter == nil {
return errFileNotOpened
} else if len(packet.Payload) == 0 {
@@ -154,11 +167,12 @@ func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error { //nolint:cyclop
}
if i.count == 0 {
i.firstFrameTimestamp = packet.Header.Timestamp
i.firstFrameTimestamp = packet.Timestamp
}
relativeTstampMs := 1000 * uint64(packet.Header.Timestamp-i.firstFrameTimestamp) / i.clockRate
relativeTstampMs := 1000 * uint64(packet.Timestamp-i.firstFrameTimestamp) / i.clockRate
if i.isVP8 { //nolint:nestif
switch i.codec {
case codecVP8:
vp8Packet := codecs.VP8Packet{}
if _, err := vp8Packet.Unmarshal(packet.Payload); err != nil {
return err
@@ -185,7 +199,35 @@ func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error { //nolint:cyclop
return err
}
i.currentFrame = nil
} else if i.isAV1 {
case codecVP9:
vp9Packet := codecs.VP9Packet{}
if _, err := vp9Packet.Unmarshal(packet.Payload); err != nil {
return err
}
switch {
case !i.seenKeyFrame && vp9Packet.P:
return nil
case i.currentFrame == nil && !vp9Packet.B:
return nil
}
i.seenKeyFrame = true
i.currentFrame = append(i.currentFrame, vp9Packet.Payload[0:]...)
if !packet.Marker {
return nil
} else if len(i.currentFrame) == 0 {
return nil
}
// the timestamp must be sequential. webrtc mandates a clock rate of 90000
// and we've assumed 30fps in the header.
if err := i.writeFrame(i.currentFrame, uint64(packet.Timestamp)/3000); err != nil {
return err
}
i.currentFrame = nil
case codecAV1:
av1Packet := &codecs.AV1Packet{}
if _, err := av1Packet.Unmarshal(packet.Payload); err != nil {
return err
@@ -201,6 +243,8 @@ func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error { //nolint:cyclop
return err
}
}
default:
return errCodecUnset
}
return nil
@@ -243,15 +287,17 @@ type Option func(i *IVFWriter) error
// WithCodec configures if IVFWriter is writing AV1 or VP8 packets to disk.
func WithCodec(mimeType string) Option {
return func(i *IVFWriter) error {
if i.isVP8 || i.isAV1 {
if i.codec != codecUnset {
return errCodecAlreadySet
}
switch mimeType {
case mimeTypeVP8:
i.isVP8 = true
i.codec = codecVP8
case mimeTypeVP9:
i.codec = codecVP9
case mimeTypeAV1:
i.isAV1 = true
i.codec = codecAV1
default:
return errNoSuchCodec
}

View File

@@ -302,3 +302,38 @@ func TestIVFWriter_AV1(t *testing.T) {
assert.NoError(t, writer.Close())
})
}
func TestIVFWriter_VP9(t *testing.T) {
buffer := &bytes.Buffer{}
writer, err := NewWith(buffer, WithCodec(mimeTypeVP9))
assert.NoError(t, err)
// No keyframe yet, ignore non-keyframe packets (P)
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0xD0, 0x02, 0xAA}}))
assert.Equal(t, buffer.Bytes(), []byte{
0x44, 0x4b, 0x49, 0x46, 0x00, 0x00, 0x20, 0x00, 0x56, 0x50, 0x39, 0x30, 0x80, 0x02, 0xe0, 0x01,
0x1e, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x84, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
})
// No current frame, ignore packets that don't start a frame (B)
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x00, 0xAA}}))
assert.Equal(t, buffer.Bytes(), []byte{
0x44, 0x4b, 0x49, 0x46, 0x00, 0x00, 0x20, 0x00, 0x56, 0x50, 0x39, 0x30, 0x80, 0x02, 0xe0, 0x01,
0x1e, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x84, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
})
// B packet, no marker bit
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x08, 0xAA}}))
assert.Equal(t, buffer.Bytes(), []byte{
0x44, 0x4b, 0x49, 0x46, 0x00, 0x00, 0x20, 0x00, 0x56, 0x50, 0x39, 0x30, 0x80, 0x02, 0xe0, 0x01,
0x1e, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x84, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
})
// B packet, Marker Bit
assert.NoError(t, writer.WriteRTP(&rtp.Packet{Header: rtp.Header{Marker: true}, Payload: []byte{0x08, 0xAB}}))
assert.Equal(t, buffer.Bytes(), []byte{
0x44, 0x4b, 0x49, 0x46, 0x00, 0x00, 0x20, 0x00, 0x56, 0x50, 0x39, 0x30, 0x80, 0x02, 0xe0, 0x01,
0x1e, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x84, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xaa, 0xab,
})
}