From 1fe21bb300b055f4b20528ffdaaa264644da3d74 Mon Sep 17 00:00:00 2001 From: Alexey Khit Date: Sat, 18 Feb 2023 19:48:02 +0300 Subject: [PATCH] Improve MPEG TS H264 processing --- pkg/h264/avc.go | 116 ++++++++++++++++++++++++++++++++++++++ pkg/mpegts/helpers.go | 83 +++------------------------ pkg/mpegts/mpegts_test.go | 34 +++++++++++ 3 files changed, 159 insertions(+), 74 deletions(-) diff --git a/pkg/h264/avc.go b/pkg/h264/avc.go index c0cdbdbd..37bf7258 100644 --- a/pkg/h264/avc.go +++ b/pkg/h264/avc.go @@ -26,6 +26,122 @@ func AnnexB2AVC(b []byte) []byte { return b } +const forbiddenZeroBit = 0x80 +const nalUnitType = 0x1F + +// DecodeStream - find and return first AU in AVC format +// useful for processing live streams with unknown separator size +func DecodeStream(annexb []byte) ([]byte, int) { + startPos := -1 + + i := 0 + for { + // search next separator + if i = IndexFrom(annexb, []byte{0, 0, 1}, i); i < 0 { + break + } + + // move i to next AU + if i += 3; i >= len(annexb) { + break + } + + // check if AU type valid + octet := annexb[i] + if octet&forbiddenZeroBit != 0 { + continue + } + + // 0 => AUD => SPS/IF/PF => AUD + // 0 => SPS/PF => SPS/PF + nalType := octet & nalUnitType + if startPos >= 0 { + switch nalType { + case NALUTypeAUD, NALUTypeSPS, NALUTypePFrame: + if annexb[i-4] == 0 { + return DecodeAnnexB(annexb[startPos : i-4]), i - 4 + } else { + return DecodeAnnexB(annexb[startPos : i-3]), i - 3 + } + } + } else { + switch nalType { + case NALUTypeSPS, NALUTypePFrame: + if i >= 4 && annexb[i-4] == 0 { + startPos = i - 4 + } else { + startPos = i - 3 + } + } + } + } + + return nil, 0 +} + +// DecodeAnnexB - convert AnnexB to AVC format +// support unknown separator size +func DecodeAnnexB(b []byte) []byte { + if b[2] == 1 { + // convert: 0 0 1 => 0 0 0 1 + b = append([]byte{0}, b...) + } + + startPos := 0 + + i := 4 + for { + // search next separato + if i = IndexFrom(b, []byte{0, 0, 1}, i); i < 0 { + break + } + + // move i to next AU + if i += 3; i >= len(b) { + break + } + + // check if AU type valid + octet := b[i] + if octet&forbiddenZeroBit != 0 { + continue + } + + switch octet & nalUnitType { + case NALUTypePFrame, NALUTypeIFrame, NALUTypeSPS, NALUTypePPS: + if b[i-4] != 0 { + // prefix: 0 0 1 + binary.BigEndian.PutUint32(b[startPos:], uint32(i-startPos-7)) + tmp := make([]byte, 0, len(b)+1) + tmp = append(tmp, b[:i]...) + tmp = append(tmp, 0) + b = append(tmp, b[i:]...) + startPos = i - 3 + } else { + // prefix: 0 0 0 1 + binary.BigEndian.PutUint32(b[startPos:], uint32(i-startPos-8)) + startPos = i - 4 + } + } + } + + binary.BigEndian.PutUint32(b[startPos:], uint32(len(b)-startPos-4)) + return b +} + +func IndexFrom(b []byte, sep []byte, from int) int { + if from > 0 { + if from < len(b) { + if i := bytes.Index(b[from:], sep); i >= 0 { + return from + i + } + } + return -1 + } + + return bytes.Index(b, sep) +} + func EncodeAVC(nals ...[]byte) (avc []byte) { var i, n int diff --git a/pkg/mpegts/helpers.go b/pkg/mpegts/helpers.go index ccf09bda..21eb7dee 100644 --- a/pkg/mpegts/helpers.go +++ b/pkg/mpegts/helpers.go @@ -1,7 +1,6 @@ package mpegts import ( - "bytes" "github.com/AlexxIT/go2rtc/pkg/h264" "github.com/AlexxIT/go2rtc/pkg/streamer" "github.com/pion/rtp" @@ -54,13 +53,8 @@ func (p *PES) SetBuffer(size uint16, b []byte) { b = b[minHeaderSize+optSize:] if p.StreamType == StreamTypeH264 { - if bytes.HasPrefix(b, []byte{0, 0, 0, 1, h264.NALUTypeAUD}) { - p.Mode = ModeStream - b = b[5:] - } - } - - if p.Mode == ModeUnknown { + p.Mode = ModeStream + } else { println("WARNING: mpegts: unknown zero-size stream") } } else { @@ -131,24 +125,23 @@ func (p *PES) GetPacket() (pkt *rtp.Packet) { p.Payload = nil case ModeStream: - i := bytes.Index(p.Payload, []byte{0, 0, 0, 1, h264.NALUTypeAUD}) - if i < 0 { - return - } - if i2 := IndexFrom(p.Payload, []byte{0, 0, 1}, i); i2 < 0 && i2 > 9 { + payload, i := h264.DecodeStream(p.Payload) + if payload == nil { return } + //log.Printf("[AVC] %v, len: %d", h264.Types(payload), len(payload)) + + p.Payload = p.Payload[i:] + pkt = &rtp.Packet{ Header: rtp.Header{ PayloadType: p.StreamType, Timestamp: uint32(time.Duration(time.Now().UnixNano()) * 90000 / time.Second), }, - Payload: DecodeAnnex3B(p.Payload[:i]), + Payload: payload, } - p.Payload = p.Payload[i+5:] - default: p.Payload = nil } @@ -191,61 +184,3 @@ func GetMedia(pkt *rtp.Packet) *streamer.Media { Codecs: []*streamer.Codec{codec}, } } - -func DecodeAnnex3B(annexb []byte) (avc []byte) { - // depends on AU delimeter size - i0 := bytes.Index(annexb, []byte{0, 0, 1}) - if i0 < 0 || i0 > 9 { - return nil - } - - annexb = annexb[i0+3:] // skip first separator - i0 = 0 - - for { - // search next separato - iN := IndexFrom(annexb, []byte{0, 0, 1}, i0) - if iN < 0 { - break - } - - // move i0 to next AU - if i0 = iN + 3; i0 >= len(annexb) { - break - } - - // check if AU type valid - octet := annexb[i0] - const forbiddenZeroBit = 0x80 - if octet&forbiddenZeroBit == 0 { - const nalUnitType = 0x1F - switch octet & nalUnitType { - case h264.NALUTypePFrame, h264.NALUTypeIFrame, h264.NALUTypeSPS, h264.NALUTypePPS: - // add AU in AVC format - avc = append(avc, byte(iN>>24), byte(iN>>16), byte(iN>>8), byte(iN)) - avc = append(avc, annexb[:iN]...) - - // cut search to next AU start - annexb = annexb[i0:] - i0 = 0 - } - } - } - - size := len(annexb) - avc = append(avc, byte(size>>24), byte(size>>16), byte(size>>8), byte(size)) - return append(avc, annexb...) -} - -func IndexFrom(b []byte, sep []byte, from int) int { - if from > 0 { - if from < len(b) { - if i := bytes.Index(b[from:], sep); i >= 0 { - return from + i - } - } - return -1 - } - - return bytes.Index(b, sep) -} diff --git a/pkg/mpegts/mpegts_test.go b/pkg/mpegts/mpegts_test.go index 840732e8..8dbf4ead 100644 --- a/pkg/mpegts/mpegts_test.go +++ b/pkg/mpegts/mpegts_test.go @@ -1,7 +1,9 @@ package mpegts import ( + "encoding/hex" "github.com/stretchr/testify/assert" + "strings" "testing" ) @@ -12,3 +14,35 @@ func TestTime(t *testing.T) { ts := ParseTime(w.Bytes()) assert.Equal(t, uint32(0xFFFFFFFF), ts) } + +func dec(s string) []byte { + s = strings.ReplaceAll(s, " ", "") + b, _ := hex.DecodeString(s) + return b +} + +func TestStream(t *testing.T) { + // ffmpeg + annexb := dec("00000001 09f0 00000001 6764001fac2484014016ec0440000003004000000c23c60c92 00000001 68ee32c8b0 000001 6588808003 00000001 09") + avc, i := ParseAVC(annexb) + assert.Equal(t, dec("00000019 6764001fac2484014016ec0440000003004000000c23c60c92 00000005 68ee32c8b0 00000005 6588808003"), avc) + assert.Equal(t, dec("00000001 09"), annexb[i:]) + + // http mpeg ts + annexb = dec("00000001 0950 000001 6764001facd2014016e8400000fa400030e081 000001 68ea8f2c 000001 65b8400eff 00000001 09") + avc, i = ParseAVC(annexb) + assert.Equal(t, dec("00000013 6764001facd2014016e8400000fa400030e081 00000004 68ea8f2c 00000005 65b8400eff"), avc) + assert.Equal(t, dec("00000001 09"), annexb[i:]) + + // tapo TC60 + annexb = dec("00000001 67640028ac1ad00a00b74dc0404050000003001000000301e8f1422a 00000001 68ee04c92240 00000001 45b80000d0 00000001 67") + avc, i = ParseAVC(annexb) + assert.Equal(t, dec("0000001C 67640028ac1ad00a00b74dc0404050000003001000000301e8f1422a 00000006 68ee04c92240 00000005 45b80000d0"), avc) + assert.Equal(t, dec("00000001 67"), annexb[i:]) + + // Tapo ? + annexb = dec("00000001 674d0032e90048014742000007d2000138d108 00000001 68ea8f20 00000001 65b8400cff 00000001 67") + avc, i = ParseAVC(annexb) + assert.Equal(t, dec("00000013 674d0032e90048014742000007d2000138d108 00000004 68ea8f20 00000005 65b8400cff"), avc) + assert.Equal(t, dec("00000001 67"), annexb[i:]) +}