pimhwang
/
mediamtx
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
Browse Source

add h264 utilities

pull/340/head
aler9 5 years ago
parent
3bf0777ee0
commit
15f0983628
6 changed files with 467 additions and 15 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 24
      internal/clientrtmp/client.go
  2. 166
      internal/h264/annexb.go
  3. 140
      internal/h264/annexb_test.go
  4. 55
      internal/h264/avcc.go
  5. 85
      internal/h264/avcc_test.go
  6. 12
      internal/sourcertmp/source.go

24
internal/clientrtmp/client.go
Unescape View File

@ -18,10 +18,11 @@ import ( @@ -18,10 +18,11 @@ import (
"github.com/aler9/gortsplib/pkg/rtpaac"
"github.com/aler9/gortsplib/pkg/rtph264"
"github.com/notedit/rtmp/av"
"github.com/notedit/rtmp/codec/h264"
rh264 "github.com/notedit/rtmp/codec/h264"
"github.com/aler9/rtsp-simple-server/internal/client"
"github.com/aler9/rtsp-simple-server/internal/externalcmd"
"github.com/aler9/rtsp-simple-server/internal/h264"
"github.com/aler9/rtsp-simple-server/internal/logger"
"github.com/aler9/rtsp-simple-server/internal/rtcpsenderset"
"github.com/aler9/rtsp-simple-server/internal/rtmputils"
@ -254,7 +255,7 @@ func (c *Client) runRead() { @@ -254,7 +255,7 @@ func (c *Client) runRead() {
c.conn.WriteMetadata(videoTrack, audioTrack)
if videoTrack != nil {
codec := h264.Codec{
codec := rh264.Codec{
SPS: map[int][]byte{
0: h264SPS,
},
@ -349,14 +350,19 @@ func (c *Client) runRead() { @@ -349,14 +350,19 @@ func (c *Client) runRead() {
// aggregate NALUs by PTS
if nt.Timestamp != videoPTS {
data, err := h264.EncodeAVCC(videoBuf)
if err != nil {
return err
}
pkt := av.Packet{
Type: av.H264,
Data: h264.FillNALUsAVCC(videoBuf),
Data: data,
Time: now.Sub(videoStartDTS),
}
c.conn.NetConn().SetWriteDeadline(time.Now().Add(c.writeTimeout))
err := c.conn.WritePacket(pkt)
err = c.conn.WritePacket(pkt)
if err != nil {
return err
}
@ -550,16 +556,16 @@ func (c *Client) runPublish() { @@ -550,16 +556,16 @@ func (c *Client) runPublish() {
return fmt.Errorf("ERR: received an H264 frame, but track is not set up")
}
// decode from AVCC format
nalus, typ := h264.SplitNALUs(pkt.Data)
if typ != h264.NALU_AVCC {
return fmt.Errorf("invalid NALU format (%d)", typ)
nalus, err := h264.DecodeAVCC(pkt.Data)
if err != nil {
return err
}
ts := pkt.Time + pkt.CTime
var nts []*rtph264.NALUAndTimestamp
for _, nt := range nalus {
nts = append(nts, &rtph264.NALUAndTimestamp{
Timestamp: pkt.Time + pkt.CTime,
Timestamp: ts,
NALU: nt,
})
}

166
internal/h264/annexb.go
Unescape View File

@ -0,0 +1,166 @@ @@ -0,0 +1,166 @@
package h264
import (
"fmt"
)
func removeAntiCompetition(nalu []byte) []byte {
// 0x00 0x00 0x03 0x00 -> 0x00 0x00 0x00
// 0x00 0x00 0x03 0x01 -> 0x00 0x00 0x01
// 0x00 0x00 0x03 0x02 -> 0x00 0x00 0x02
// 0x00 0x00 0x03 0x03 -> 0x00 0x00 0x03
var ret []byte
step := 0
start := 0
for i, b := range nalu {
switch step {
case 0:
if b == 0 {
step++
}
case 1:
if b == 0 {
step++
} else {
step = 0
}
case 2:
if b == 3 {
step++
} else {
step = 0
}
case 3:
switch b {
case 3, 2, 1, 0:
ret = append(ret, nalu[start:i-3]...)
ret = append(ret, []byte{0x00, 0x00, b}...)
step = 0
start = i + 1
default:
step = 0
}
}
}
ret = append(ret, nalu[start:]...)
return ret
}
func addAntiCompetition(dest []byte, nalu []byte) []byte {
step := 0
start := 0
for i, b := range nalu {
switch step {
case 0:
if b == 0 {
step++
}
case 1:
if b == 0 {
step++
} else {
step = 0
}
case 2:
switch b {
case 3, 2, 1, 0:
dest = append(dest, nalu[start:i-2]...)
dest = append(dest, []byte{0x00, 0x00, 0x03, b}...)
step = 0
start = i + 1
default:
step = 0
}
}
}
dest = append(dest, nalu[start:]...)
return dest
}
// DecodeAnnexB decodes NALUs from the Annex-B code stream format.
func DecodeAnnexB(byts []byte) ([][]byte, error) {
bl := len(byts)
// check initial delimiter
n := func() int {
if bl < 3 || byts[0] != 0x00 || byts[1] != 0x00 {
return -1
}
if byts[2] == 0x01 {
return 3
}
if bl < 4 || byts[2] != 0x00 || byts[3] != 0x01 {
return -1
}
return 4
}()
if n < 0 {
return nil, fmt.Errorf("input doesn't start with a delimiter")
}
var ret [][]byte
zeros := 0
start := n
delimStart := 0
for i := n; i < bl; i++ {
switch byts[i] {
case 0:
if zeros == 0 {
delimStart = i
}
zeros++
case 1:
if zeros == 2 || zeros == 3 {
nalu := byts[start:delimStart]
if len(nalu) == 0 {
return nil, fmt.Errorf("empty NALU")
}
ret = append(ret, removeAntiCompetition(nalu))
start = i + 1
}
zeros = 0
default:
zeros = 0
}
}
nalu := byts[start:bl]
if len(nalu) == 0 {
return nil, fmt.Errorf("empty NALU")
}
ret = append(ret, removeAntiCompetition(nalu))
return ret, nil
}
// EncodeAnnexB encodes NALUs into the Annex-B code stream format.
func EncodeAnnexB(nalus [][]byte) ([]byte, error) {
var ret []byte
for _, nalu := range nalus {
ret = append(ret, []byte{0x00, 0x00, 0x00, 0x01}...)
ret = addAntiCompetition(ret, nalu)
}
return ret, nil
}

140
internal/h264/annexb_test.go
Unescape View File

@ -0,0 +1,140 @@ @@ -0,0 +1,140 @@
package h264
import (
"testing"
"github.com/stretchr/testify/require"
)
var annexBCases = []struct {
name string
encin []byte
encout []byte
dec [][]byte
}{
{
"2 zeros, single",
[]byte{0x00, 0x00, 0x01, 0xaa, 0xbb},
[]byte{0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb},
[][]byte{
{0xaa, 0xbb},
},
},
{
"2 zeros, multiple",
[]byte{
0x00, 0x00, 0x01, 0xaa, 0xbb, 0x00, 0x00, 0x01,
0xcc, 0xdd, 0x00, 0x00, 0x01, 0xee, 0xff,
},
[]byte{
0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb, 0x00, 0x00,
0x00, 0x01, 0xcc, 0xdd, 0x00, 0x00, 0x00, 0x01,
0xee, 0xff,
},
[][]byte{
{0xaa, 0xbb},
{0xcc, 0xdd},
{0xee, 0xff},
},
},
{
"3 zeros, single",
[]byte{0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb},
[]byte{0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb},
[][]byte{
{0xaa, 0xbb},
},
},
{
"3 zeros, multiple",
[]byte{
0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb, 0x00, 0x00,
0x00, 0x01, 0xcc, 0xdd, 0x00, 0x00, 0x00, 0x01,
0xee, 0xff,
},
[]byte{
0x00, 0x00, 0x00, 0x01, 0xaa, 0xbb, 0x00, 0x00,
0x00, 0x01, 0xcc, 0xdd, 0x00, 0x00, 0x00, 0x01,
0xee, 0xff,
},
[][]byte{
{0xaa, 0xbb},
{0xcc, 0xdd},
{0xee, 0xff},
},
},
{
"anti-competition",
[]byte{
0x00, 0x00, 0x00, 0x01,
0x00, 0x00, 0x03, 0x00,
0x00, 0x00, 0x03, 0x01,
0x00, 0x00, 0x03, 0x02,
0x00, 0x00, 0x03, 0x03,
},
[]byte{
0x00, 0x00, 0x00, 0x01,
0x00, 0x00, 0x03, 0x00,
0x00, 0x00, 0x03, 0x01,
0x00, 0x00, 0x03, 0x02,
0x00, 0x00, 0x03, 0x03,
},
[][]byte{
{
0x00, 0x00, 0x00,
0x00, 0x00, 0x01,
0x00, 0x00, 0x02,
0x00, 0x00, 0x03,
},
},
},
}
func TestAnnexBDecode(t *testing.T) {
for _, ca := range annexBCases {
t.Run(ca.name, func(t *testing.T) {
dec, err := DecodeAnnexB(ca.encin)
require.NoError(t, err)
require.Equal(t, ca.dec, dec)
})
}
}
func TestAnnexBEncode(t *testing.T) {
for _, ca := range annexBCases {
t.Run(ca.name, func(t *testing.T) {
enc, err := EncodeAnnexB(ca.dec)
require.NoError(t, err)
require.Equal(t, ca.encout, enc)
})
}
}
func TestAnnexBDecodeError(t *testing.T) {
for _, ca := range []struct {
name string
enc []byte
}{
{
"empty",
[]byte{},
},
{
"missing initial delimiter",
[]byte{0xaa, 0xbb},
},
{
"empty initial",
[]byte{0x00, 0x00, 0x01},
},
{
"empty 2nd",
[]byte{0x00, 0x00, 0x01, 0xaa, 0x00, 0x00, 0x01},
},
} {
t.Run(ca.name, func(t *testing.T) {
_, err := DecodeAnnexB(ca.enc)
require.Error(t, err)
})
}
}

55
internal/h264/avcc.go
Unescape View File

@ -0,0 +1,55 @@ @@ -0,0 +1,55 @@
package h264
import (
"encoding/binary"
"fmt"
)
// DecodeAVCC encodes NALUs from the AVCC code stream format.
func DecodeAVCC(byts []byte) ([][]byte, error) {
var ret [][]byte
for len(byts) > 0 {
if len(byts) < 4 {
return nil, fmt.Errorf("invalid length")
}
le := binary.BigEndian.Uint32(byts)
byts = byts[4:]
if len(byts) < int(le) {
return nil, fmt.Errorf("invalid length")
}
ret = append(ret, byts[:le])
byts = byts[le:]
}
if len(ret) == 0 {
return nil, fmt.Errorf("no NALUs decoded")
}
return ret, nil
}
// EncodeAVCC encodes NALUs into the AVCC code stream format.
func EncodeAVCC(nalus [][]byte) ([]byte, error) {
le := 0
for _, nalu := range nalus {
le += 4 + len(nalu)
}
ret := make([]byte, le)
pos := 0
for _, nalu := range nalus {
ln := len(nalu)
binary.BigEndian.PutUint32(ret[pos:], uint32(ln))
pos += 4
copy(ret[pos:], nalu)
pos += ln
}
return ret, nil
}

85
internal/h264/avcc_test.go
Unescape View File

@ -0,0 +1,85 @@ @@ -0,0 +1,85 @@
package h264
import (
"testing"
"github.com/stretchr/testify/require"
)
var casesAVCC = []struct {
name string
enc []byte
dec [][]byte
}{
{
"single",
[]byte{
0x00, 0x00, 0x00, 0x03,
0xaa, 0xbb, 0xcc,
},
[][]byte{
{0xaa, 0xbb, 0xcc},
},
},
{
"multiple",
[]byte{
0x00, 0x00, 0x00, 0x02,
0xaa, 0xbb,
0x00, 0x00, 0x00, 0x02,
0xcc, 0xdd,
0x00, 0x00, 0x00, 0x02,
0xee, 0xff,
},
[][]byte{
{0xaa, 0xbb},
{0xcc, 0xdd},
{0xee, 0xff},
},
},
}
func TestAVCCDecode(t *testing.T) {
for _, ca := range casesAVCC {
t.Run(ca.name, func(t *testing.T) {
dec, err := DecodeAVCC(ca.enc)
require.NoError(t, err)
require.Equal(t, ca.dec, dec)
})
}
}
func TestAVCCEncode(t *testing.T) {
for _, ca := range casesAVCC {
t.Run(ca.name, func(t *testing.T) {
enc, err := EncodeAVCC(ca.dec)
require.NoError(t, err)
require.Equal(t, ca.enc, enc)
})
}
}
func TestAVCCDecodeError(t *testing.T) {
for _, ca := range []struct {
name string
enc []byte
}{
{
"empty",
[]byte{},
},
{
"invalid length",
[]byte{0x01},
},
{
"invalid length",
[]byte{0x00, 0x00, 0x00, 0x03},
},
} {
t.Run(ca.name, func(t *testing.T) {
_, err := DecodeAVCC(ca.enc)
require.Error(t, err)
})
}
}

12
internal/sourcertmp/source.go
Unescape View File

@ -11,9 +11,9 @@ import ( @@ -11,9 +11,9 @@ import (
"github.com/aler9/gortsplib/pkg/rtpaac"
"github.com/aler9/gortsplib/pkg/rtph264"
"github.com/notedit/rtmp/av"
"github.com/notedit/rtmp/codec/h264"
"github.com/notedit/rtmp/format/rtmp"
"github.com/aler9/rtsp-simple-server/internal/h264"
"github.com/aler9/rtsp-simple-server/internal/logger"
"github.com/aler9/rtsp-simple-server/internal/rtcpsenderset"
"github.com/aler9/rtsp-simple-server/internal/rtmputils"
@ -214,16 +214,16 @@ func (s *Source) runInner() bool { @@ -214,16 +214,16 @@ func (s *Source) runInner() bool {
return fmt.Errorf("ERR: received an H264 frame, but track is not set up")
}
// decode from AVCC format
nalus, typ := h264.SplitNALUs(pkt.Data)
if typ != h264.NALU_AVCC {
return fmt.Errorf("invalid NALU format (%d)", typ)
nalus, err := h264.DecodeAVCC(pkt.Data)
if err != nil {
return err
}
ts := pkt.Time + pkt.CTime
var nts []*rtph264.NALUAndTimestamp
for _, nt := range nalus {
nts = append(nts, &rtph264.NALUAndTimestamp{
Timestamp: pkt.Time + pkt.CTime,
Timestamp: ts,
NALU: nt,
})
}

Write Preview
Loading…
Cancel
Save