Move IVFWriter to pkg

Allow this to be used/improved by others. SampleBuilder will be a
sibling

Resolves #115
This commit is contained in:
Sean DuBois
2018-08-27 00:29:48 -07:00
parent fde197b31c
commit 2adb8dcc80
2 changed files with 25 additions and 20 deletions

View File

@@ -9,6 +9,7 @@ import (
"github.com/pions/webrtc" "github.com/pions/webrtc"
"github.com/pions/webrtc/pkg/ice" "github.com/pions/webrtc/pkg/ice"
"github.com/pions/webrtc/pkg/media/ivfwriter"
) )
func main() { func main() {
@@ -49,12 +50,14 @@ func main() {
peerConnection.Ontrack = func(track *webrtc.RTCTrack) { peerConnection.Ontrack = func(track *webrtc.RTCTrack) {
if track.Codec.Name == webrtc.VP8 { if track.Codec.Name == webrtc.VP8 {
fmt.Println("Got VP8 track, saving to disk as output.ivf") fmt.Println("Got VP8 track, saving to disk as output.ivf")
i, err := newIVFWriter("output.ivf") i, err := ivfwriter.New("output.ivf")
if err != nil { if err != nil {
panic(err) panic(err)
} }
for { for {
i.addPacket(<-track.Packets) if err := i.AddPacket(<-track.Packets); err != nil {
panic(err)
}
} }
} }
} }

View File

@@ -1,4 +1,4 @@
package main package ivfwriter
import ( import (
"encoding/binary" "encoding/binary"
@@ -9,19 +9,15 @@ import (
"github.com/pions/webrtc/pkg/rtp/codecs" "github.com/pions/webrtc/pkg/rtp/codecs"
) )
type ivfWriter struct { // IVFWriter is used to take RTP packets and write them to an IVF on disk
type IVFWriter struct {
fd *os.File fd *os.File
count uint64 count uint64
currentFrame []byte currentFrame []byte
} }
func panicWrite(fd *os.File, data []byte) { // New builds a new IVF writer
if _, err := fd.Write(data); err != nil { func New(fileName string) (*IVFWriter, error) {
panic(err)
}
}
func newIVFWriter(fileName string) (*ivfWriter, error) {
f, err := os.Create(fileName) f, err := os.Create(fileName)
if err != nil { if err != nil {
return nil, err return nil, err
@@ -39,27 +35,29 @@ func newIVFWriter(fileName string) (*ivfWriter, error) {
binary.LittleEndian.PutUint32(header[24:], 900) // Frame count binary.LittleEndian.PutUint32(header[24:], 900) // Frame count
binary.LittleEndian.PutUint32(header[28:], 0) // Unused binary.LittleEndian.PutUint32(header[28:], 0) // Unused
panicWrite(f, header) if _, err := f.Write(header); err != nil {
return nil, err
}
i := &ivfWriter{fd: f} return &IVFWriter{fd: f}, nil
return i, nil
} }
func (i *ivfWriter) addPacket(packet *rtp.Packet) { // AddPacket adds a new packet and writes the appropriate headers for it
func (i *IVFWriter) AddPacket(packet *rtp.Packet) error {
vp8Packet := codecs.VP8Packet{} vp8Packet := codecs.VP8Packet{}
err := vp8Packet.Unmarshal(packet) err := vp8Packet.Unmarshal(packet)
if err != nil { if err != nil {
panic(err) return err
} }
i.currentFrame = append(i.currentFrame, vp8Packet.Payload[0:]...) i.currentFrame = append(i.currentFrame, vp8Packet.Payload[0:]...)
if !packet.Marker { if !packet.Marker {
return return nil
} else if len(i.currentFrame) == 0 { } else if len(i.currentFrame) == 0 {
fmt.Println("skipping") fmt.Println("skipping")
return return nil
} }
frameHeader := make([]byte, 12) frameHeader := make([]byte, 12)
@@ -68,8 +66,12 @@ func (i *ivfWriter) addPacket(packet *rtp.Packet) {
i.count++ i.count++
panicWrite(i.fd, frameHeader) if _, err := i.fd.Write(frameHeader); err != nil {
panicWrite(i.fd, i.currentFrame) return err
} else if _, err := i.fd.Write(i.currentFrame); err != nil {
return err
}
i.currentFrame = nil i.currentFrame = nil
return nil
} }