mirror of
https://github.com/bluenviron/mediamtx.git
synced 2025-12-25 20:41:59 -08:00
add h264 utilities
This commit is contained in:
aler9
parent
3bf0777ee0
commit
15f0983628
6 changed files with 467 additions and 15 deletions
|
|
@ -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() {
|
|||
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() {
|
|||
|
||||
// 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() {
|
|||
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
Normal file
166
internal/h264/annexb.go
Normal file
|
|
@ -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
Normal file
140
internal/h264/annexb_test.go
Normal file
|
|
@ -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
Normal file
55
internal/h264/avcc.go
Normal file
|
|
@ -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
Normal file
85
internal/h264/avcc_test.go
Normal file
|
|
@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -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 {
|
|||
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,
|
||||
})
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue