Improve MPEG TS H264 processing

This commit is contained in:
Alexey Khit
2023-02-18 19:48:02 +03:00
parent 41cdcb69c6
commit 1fe21bb300
3 changed files with 159 additions and 74 deletions
+116
View File
@@ -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
+9 -74
View File
@@ -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)
}
+34
View File
@@ -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:])
}