From 7d29ef5a64162430f34e42017bfb23206cd9f107 Mon Sep 17 00:00:00 2001 From: Juliusz Chroboczek Date: Thu, 29 Jul 2021 23:01:22 +0200 Subject: [PATCH] Make disk writer use the codecs package. Instead of testing whether a sample is a keyframe, we test at the packet level, then compare timestamps to identify the keyframe. --- codecs/codecs.go | 46 +++++++++++++++++ diskwriter/diskwriter.go | 104 +++++++-------------------------------- 2 files changed, 63 insertions(+), 87 deletions(-) diff --git a/codecs/codecs.go b/codecs/codecs.go index 6ffbec3..7bb8e98 100644 --- a/codecs/codecs.go +++ b/codecs/codecs.go @@ -182,6 +182,52 @@ func Keyframe(codec string, packet *rtp.Packet) (bool, bool) { return false, false } +func KeyframeDimensions(codec string, packet *rtp.Packet) (uint32, uint32) { + if strings.EqualFold(codec, "video/vp8") { + var vp8 codecs.VP8Packet + _, err := vp8.Unmarshal(packet.Payload) + if err != nil { + return 0, 0 + } + if len(vp8.Payload) < 10 { + return 0, 0 + } + raw := uint32(vp8.Payload[6]) | uint32(vp8.Payload[7])<<8 | + uint32(vp8.Payload[8])<<16 | uint32(vp8.Payload[9])<<24 + width := raw & 0x3FFF + height := (raw >> 16) & 0x3FFF + return width, height + } else if strings.EqualFold(codec, "video/vp9") { + if packet == nil { + return 0, 0 + } + var vp9 codecs.VP9Packet + _, err := vp9.Unmarshal(packet.Payload) + if err != nil { + return 0, 0 + } + if !vp9.V { + return 0, 0 + } + w := uint32(0) + h := uint32(0) + for i := range vp9.Width { + if i >= len(vp9.Height) { + break + } + if w < uint32(vp9.Width[i]) { + w = uint32(vp9.Width[i]) + } + if h < uint32(vp9.Height[i]) { + h = uint32(vp9.Height[i]) + } + } + return w, h + } else { + return 0, 0 + } +} + type Flags struct { Seqno uint16 Start bool diff --git a/diskwriter/diskwriter.go b/diskwriter/diskwriter.go index f3e084f..0853ad9 100644 --- a/diskwriter/diskwriter.go +++ b/diskwriter/diskwriter.go @@ -20,6 +20,7 @@ import ( "github.com/jech/samplebuilder" + gcodecs "github.com/jech/galene/codecs" "github.com/jech/galene/conn" "github.com/jech/galene/group" ) @@ -411,70 +412,6 @@ func (t *diskTrack) SetTimeOffset(ntp uint64, rtp uint32) { func (t *diskTrack) SetCname(string) { } -func isKeyframe(codec string, data []byte) bool { - if strings.EqualFold(codec, "video/vp8") { - if len(data) < 1 { - return false - } - return (data[0] & 0x1) == 0 - } else if strings.EqualFold(codec, "video/vp9") { - if len(data) < 1 { - return false - } - if data[0]&0xC0 != 0x80 { - return false - } - profile := (data[0] >> 4) & 0x3 - if profile != 3 { - return (data[0] & 0xC) == 0 - } - return (data[0] & 0x6) == 0 - } else { - panic("Eek!") - } -} - -func keyframeDimensions(codec string, data []byte, packet *rtp.Packet) (uint32, uint32) { - if strings.EqualFold(codec, "video/vp8") { - if len(data) < 10 { - return 0, 0 - } - raw := uint32(data[6]) | uint32(data[7])<<8 | - uint32(data[8])<<16 | uint32(data[9])<<24 - width := raw & 0x3FFF - height := (raw >> 16) & 0x3FFF - return width, height - } else if strings.EqualFold(codec, "video/vp9") { - if packet == nil { - return 0, 0 - } - var vp9 codecs.VP9Packet - _, err := vp9.Unmarshal(packet.Payload) - if err != nil { - return 0, 0 - } - if !vp9.V { - return 0, 0 - } - w := uint32(0) - h := uint32(0) - for i := range vp9.Width { - if i >= len(vp9.Height) { - break - } - if w < uint32(vp9.Width[i]) { - w = uint32(vp9.Width[i]) - } - if h < uint32(vp9.Height[i]) { - h = uint32(vp9.Height[i]) - } - } - return w, h - } else { - return 0, 0 - } -} - func (t *diskTrack) Write(buf []byte) (int, error) { t.conn.mu.Lock() defer t.conn.mu.Unlock() @@ -527,21 +464,11 @@ func (t *diskTrack) Write(buf []byte) (int, error) { // writeRTP writes the packet without doing any loss recovery. // Called locked. func (t *diskTrack) writeRTP(p *rtp.Packet) error { - codec := t.remote.Codec() - if strings.EqualFold(codec.MimeType, "video/vp9") { - var vp9 codecs.VP9Packet - _, err := vp9.Unmarshal(p.Payload) - if err == nil && vp9.B && len(vp9.Payload) >= 1 { - profile := (vp9.Payload[0] >> 4) & 0x3 - kf := false - if profile != 3 { - kf = (vp9.Payload[0] & 0xC) == 0 - } else { - kf = (vp9.Payload[0] & 0x6) == 0 - } - if kf { - t.savedKf = p - } + codec := t.remote.Codec().MimeType + if len(codec) > 6 && strings.EqualFold(codec[:6], "video/") { + kf, _ := gcodecs.Keyframe(codec, p) + if kf { + t.savedKf = p } } @@ -554,7 +481,7 @@ func (t *diskTrack) writeRTP(p *rtp.Packet) error { // then samples will be flushed even if they are preceded by incomplete // samples. func (t *diskTrack) writeBuffered(force bool) error { - codec := t.remote.Codec() + codec := t.remote.Codec().MimeType for { var sample *media.Sample @@ -568,16 +495,18 @@ func (t *diskTrack) writeBuffered(force bool) error { return nil } - keyframe := true + var keyframe bool + if len(codec) > 6 && strings.EqualFold(codec[:6], "video/") { + if t.savedKf == nil { + keyframe = false + } else { + keyframe = (ts == t.savedKf.Timestamp) + } - if strings.EqualFold(codec.MimeType, "video/vp8") || - strings.EqualFold(codec.MimeType, "video/vp9") { - keyframe = isKeyframe(codec.MimeType, sample.Data) if keyframe { err := t.conn.initWriter( - keyframeDimensions( - codec.MimeType, sample.Data, - t.savedKf, + gcodecs.KeyframeDimensions( + codec, t.savedKf, ), ) if err != nil { @@ -588,6 +517,7 @@ func (t *diskTrack) writeBuffered(force bool) error { } } } else { + keyframe = true if t.writer == nil { if !t.conn.hasVideo { err := t.conn.initWriter(0, 0)