mirror of
https://github.com/pion/webrtc.git
synced 2025-10-05 15:16:52 +08:00
137 lines
3.2 KiB
Go
137 lines
3.2 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"fmt"
|
|
"os"
|
|
"time"
|
|
|
|
"github.com/pions/webrtc/rtp"
|
|
)
|
|
|
|
type IVFWriter struct {
|
|
fd *os.File
|
|
time time.Time
|
|
count uint64
|
|
currentFrame []byte
|
|
}
|
|
|
|
type VP8RTPPacket struct {
|
|
// Required Header
|
|
X uint8 /* extended controlbits present */
|
|
N uint8 /* (non-reference frame) when set to 1 this frame can be discarded */
|
|
S uint8 /* start of VP8 partition */
|
|
PID uint8 /* partition index */
|
|
|
|
// Optional Header
|
|
I uint8 /* 1 if PictureID is present */
|
|
L uint8 /* 1 if TL0PICIDX is present */
|
|
T uint8 /* 1 if TID is present */
|
|
K uint8 /* 1 if KEYIDX is present */
|
|
PictureID uint16 /* 8 or 16 bits, picture ID */
|
|
TL0PICIDX uint8 /* 8 bits temporal level zero index */
|
|
|
|
Payload []byte
|
|
}
|
|
|
|
func panicWrite(fd *os.File, data []byte) {
|
|
if _, err := fd.Write(data); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
func NewIVFWriter(fileName string) (*IVFWriter, error) {
|
|
f, err := os.Create(fileName)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
header := make([]byte, 32)
|
|
copy(header[0:], []byte("DKIF")) // DKIF
|
|
binary.LittleEndian.PutUint16(header[4:], 0) // Version
|
|
binary.LittleEndian.PutUint16(header[6:], 32) // Header Size
|
|
copy(header[8:], []byte("VP80")) // FOURCC
|
|
binary.LittleEndian.PutUint16(header[12:], 640) // Version
|
|
binary.LittleEndian.PutUint16(header[14:], 480) // Header Size
|
|
binary.LittleEndian.PutUint32(header[16:], 30) // Framerate numerator
|
|
binary.LittleEndian.PutUint32(header[20:], 1) // Framerate Denominator
|
|
binary.LittleEndian.PutUint32(header[24:], 900) // Frame count
|
|
binary.LittleEndian.PutUint32(header[28:], 0) // Unused
|
|
|
|
panicWrite(f, header)
|
|
|
|
i := &IVFWriter{fd: f}
|
|
return i, nil
|
|
}
|
|
|
|
func (i *IVFWriter) DecodeVP8RTPPacket(packet *rtp.Packet) (*VP8RTPPacket, error) {
|
|
p := packet.Payload
|
|
|
|
vp8Packet := &VP8RTPPacket{}
|
|
|
|
payloadIndex := 0
|
|
vp8Packet.X = (p[payloadIndex] & 0x80) >> 7
|
|
vp8Packet.N = (p[payloadIndex] & 0x20) >> 5
|
|
vp8Packet.S = (p[payloadIndex] & 0x10) >> 4
|
|
vp8Packet.PID = p[payloadIndex] & 0x07
|
|
|
|
payloadIndex++
|
|
|
|
if vp8Packet.X == 1 {
|
|
vp8Packet.I = (p[payloadIndex] & 0x80) >> 7
|
|
vp8Packet.L = (p[payloadIndex] & 0x40) >> 6
|
|
vp8Packet.T = (p[payloadIndex] & 0x20) >> 5
|
|
vp8Packet.K = (p[payloadIndex] & 0x10) >> 4
|
|
payloadIndex++
|
|
}
|
|
|
|
if vp8Packet.I == 1 { // PID present?
|
|
if p[payloadIndex]&0x80 > 0 { // M == 1, PID is 16bit
|
|
payloadIndex += 2
|
|
} else {
|
|
payloadIndex++
|
|
}
|
|
}
|
|
|
|
if vp8Packet.L == 1 {
|
|
payloadIndex++
|
|
}
|
|
|
|
if vp8Packet.T == 1 || vp8Packet.K == 1 {
|
|
payloadIndex++
|
|
}
|
|
|
|
vp8Packet.Payload = p[payloadIndex:]
|
|
|
|
return vp8Packet, nil
|
|
}
|
|
|
|
func (i *IVFWriter) AddPacket(packet *rtp.Packet) {
|
|
|
|
vp8Packet, _ := i.DecodeVP8RTPPacket(packet)
|
|
|
|
i.currentFrame = append(i.currentFrame, vp8Packet.Payload[0:]...)
|
|
|
|
if !packet.Marker {
|
|
return
|
|
} else if len(i.currentFrame) == 0 {
|
|
fmt.Println("skipping")
|
|
return
|
|
}
|
|
|
|
frameHeader := make([]byte, 12)
|
|
binary.LittleEndian.PutUint32(frameHeader[0:], uint32(len(i.currentFrame))) // Frame length
|
|
binary.LittleEndian.PutUint64(frameHeader[4:], i.count) // PTS
|
|
|
|
i.count += 1
|
|
|
|
panicWrite(i.fd, frameHeader)
|
|
panicWrite(i.fd, i.currentFrame)
|
|
|
|
i.currentFrame = nil
|
|
}
|
|
|
|
func (i *IVFWriter) Close() error {
|
|
return i.fd.Close()
|
|
}
|