mirror of
https://github.com/Monibuca/engine.git
synced 2025-11-01 20:32:44 +08:00
411 lines
10 KiB
Go
411 lines
10 KiB
Go
package engine
|
||
|
||
import (
|
||
"container/ring"
|
||
"time"
|
||
|
||
"github.com/Monibuca/utils/v3"
|
||
"github.com/Monibuca/utils/v3/codec"
|
||
)
|
||
|
||
const (
|
||
naluTypeBitmask = 0b0001_1111
|
||
naluTypeBitmask_hevc = 0x7E
|
||
)
|
||
|
||
type VideoPack struct {
|
||
AVPack
|
||
CompositionTime uint32
|
||
NALUs [][]byte
|
||
IDR bool // 是否关键帧
|
||
}
|
||
|
||
func (v *VideoPack) ResetNALUs() {
|
||
if cap(v.NALUs) > 0 {
|
||
v.NALUs = v.NALUs[:0]
|
||
}
|
||
}
|
||
|
||
func (v *VideoPack) SetNalu0(nalu []byte) {
|
||
if cap(v.NALUs) > 0 {
|
||
v.NALUs = v.NALUs[:1]
|
||
v.NALUs[0] = nalu
|
||
} else {
|
||
v.NALUs = [][]byte{nalu}
|
||
}
|
||
}
|
||
|
||
type VideoTrack struct {
|
||
IDRing *ring.Ring `json:"-"` //最近的关键帧位置，首屏渲染
|
||
AVTrack
|
||
SPSInfo codec.SPSInfo
|
||
GOP int //关键帧间隔
|
||
ExtraData *VideoPack `json:"-"` //H264(SPS、PPS) H265(VPS、SPS、PPS)
|
||
WaitIDR chan struct{} `json:"-"`
|
||
revIDR func()
|
||
PushNalu func(ts uint32, cts uint32, nalus ...[]byte) `json:"-"`
|
||
UsingDonlField bool
|
||
writeByteStream func()
|
||
idrCount int //处于缓冲中的关键帧数量
|
||
nalulenSize int
|
||
*VideoPack `json:"-"` //当前写入的视频数据
|
||
}
|
||
|
||
func (s *Stream) NewVideoTrack(codec byte) (vt *VideoTrack) {
|
||
vt = &VideoTrack{
|
||
WaitIDR: make(chan struct{}),
|
||
revIDR: func() {
|
||
vt.IDRing = vt.Ring
|
||
close(vt.WaitIDR)
|
||
idrSequence := vt.Sequence
|
||
vt.ts = vt.Timestamp
|
||
vt.idrCount++
|
||
vt.revIDR = func() {
|
||
vt.idrCount++
|
||
vt.GOP = vt.Sequence - idrSequence
|
||
if l := vt.Ring.Len() - vt.GOP - 5; l > 5 {
|
||
//缩小缓冲环节省内存
|
||
vt.Unlink(l).Do(func(v interface{}) {
|
||
if v.(*AVItem).Value.(*VideoPack).IDR {
|
||
vt.idrCount--
|
||
}
|
||
})
|
||
}
|
||
vt.IDRing = vt.Ring
|
||
idrSequence = vt.Sequence
|
||
vt.resetBPS()
|
||
}
|
||
},
|
||
}
|
||
vt.PushNalu = vt.pushNalu
|
||
vt.Stream = s
|
||
vt.CodecID = codec
|
||
vt.Init(s.Context, 256)
|
||
vt.poll = time.Millisecond * 20
|
||
vt.Do(func(v interface{}) {
|
||
v.(*AVItem).Value = new(VideoPack)
|
||
})
|
||
vt.setCurrent()
|
||
switch codec {
|
||
case 7:
|
||
s.VideoTracks.AddTrack("h264", vt)
|
||
case 12:
|
||
s.VideoTracks.AddTrack("h265", vt)
|
||
}
|
||
return
|
||
}
|
||
|
||
func (vt *VideoTrack) PushAnnexB(ts uint32, cts uint32, payload []byte) {
|
||
vt.PushNalu(ts, cts, codec.SplitH264(payload)...)
|
||
}
|
||
|
||
func (vt *VideoTrack) pushNalu(ts uint32, cts uint32, nalus ...[]byte) {
|
||
idrBit := 0x10 | vt.CodecID
|
||
nIdrBit := 0x20 | vt.CodecID
|
||
tmp := make([]byte, 4)
|
||
// 缓冲中只包含Nalu数据所以写入rtmp格式时需要按照ByteStream格式写入
|
||
vt.writeByteStream = func() {
|
||
vt.Reset()
|
||
if vt.IDR {
|
||
tmp[0] = idrBit
|
||
} else {
|
||
tmp[0] = nIdrBit
|
||
}
|
||
tmp[1] = 1
|
||
vt.Buffer.Write(tmp[:2])
|
||
utils.BigEndian.PutUint24(tmp, vt.CompositionTime)
|
||
vt.Buffer.Write(tmp[:3])
|
||
for _, nalu := range vt.NALUs {
|
||
utils.BigEndian.PutUint32(tmp, uint32(len(nalu)))
|
||
vt.Buffer.Write(tmp)
|
||
vt.Buffer.Write(nalu)
|
||
}
|
||
vt.Bytes2Payload()
|
||
}
|
||
switch vt.CodecID {
|
||
case 7:
|
||
{
|
||
var info codec.AVCDecoderConfigurationRecord
|
||
vt.PushNalu = func(ts uint32, cts uint32, nalus ...[]byte) {
|
||
// 等待接收SPS和PPS数据
|
||
for _, nalu := range nalus {
|
||
if len(nalu) == 0 {
|
||
continue
|
||
}
|
||
switch nalu[0] & naluTypeBitmask {
|
||
case codec.NALU_SPS:
|
||
info.SequenceParameterSetNALUnit = nalu
|
||
info.SequenceParameterSetLength = uint16(len(nalu))
|
||
vt.SPSInfo, _ = codec.ParseSPS(nalu)
|
||
case codec.NALU_PPS:
|
||
info.PictureParameterSetNALUnit = nalu
|
||
info.PictureParameterSetLength = uint16(len(nalu))
|
||
}
|
||
}
|
||
if info.SequenceParameterSetNALUnit != nil && info.PictureParameterSetNALUnit != nil {
|
||
vt.ExtraData = &VideoPack{
|
||
NALUs: [][]byte{info.SequenceParameterSetNALUnit, info.PictureParameterSetNALUnit},
|
||
}
|
||
vt.ExtraData.Payload = codec.BuildH264SeqHeaderFromSpsPps(info.SequenceParameterSetNALUnit, info.PictureParameterSetNALUnit)
|
||
}
|
||
if vt.ExtraData == nil {
|
||
return
|
||
}
|
||
//已完成SPS和PPS 组装，重置push函数，接收视频数据
|
||
vt.PushNalu = func(ts uint32, cts uint32, nalus ...[]byte) {
|
||
var nonIDRs int
|
||
for _, nalu := range nalus {
|
||
naluLen := len(nalu)
|
||
if naluLen == 0 {
|
||
continue
|
||
}
|
||
naluType := nalu[0] & naluTypeBitmask
|
||
switch naluType {
|
||
case codec.NALU_SPS:
|
||
vt.ExtraData.NALUs[0] = nalu
|
||
vt.SPSInfo, _ = codec.ParseSPS(nalu)
|
||
case codec.NALU_PPS:
|
||
vt.ExtraData.NALUs[1] = nalu
|
||
vt.ExtraData.Payload = codec.BuildH264SeqHeaderFromSpsPps(vt.ExtraData.NALUs[0], vt.ExtraData.NALUs[1])
|
||
case codec.NALU_Access_Unit_Delimiter:
|
||
case codec.NALU_IDR_Picture:
|
||
if nonIDRs > 0 {
|
||
vt.push()
|
||
nonIDRs = 0
|
||
}
|
||
vt.addBytes(naluLen)
|
||
vt.IDR = true
|
||
vt.SetTs(ts)
|
||
vt.CompositionTime = cts
|
||
vt.SetNalu0(nalu)
|
||
vt.push()
|
||
case codec.NALU_Non_IDR_Picture:
|
||
vt.addBytes(naluLen)
|
||
vt.IDR = false
|
||
vt.SetTs(ts)
|
||
vt.CompositionTime = cts
|
||
if nonIDRs == 0 {
|
||
vt.SetNalu0(nalu)
|
||
} else {
|
||
vt.NALUs = append(vt.NALUs, nalu)
|
||
}
|
||
nonIDRs++
|
||
case codec.NALU_SEI:
|
||
case codec.NALU_Filler_Data:
|
||
default:
|
||
utils.Printf("%s,nalType not support yet:%d,[0]=0x%X", vt.Stream.StreamPath, naluType, nalu[0])
|
||
}
|
||
}
|
||
if nonIDRs > 0 {
|
||
vt.push()
|
||
}
|
||
}
|
||
}
|
||
}
|
||
case 12:
|
||
var vps, sps, pps []byte
|
||
vt.PushNalu = func(ts uint32, cts uint32, nalus ...[]byte) {
|
||
// 等待接收SPS和PPS数据
|
||
for _, nalu := range nalus {
|
||
if len(nalu) == 0 {
|
||
continue
|
||
}
|
||
switch nalu[0] & naluTypeBitmask_hevc >> 1 {
|
||
case codec.NAL_UNIT_VPS:
|
||
vps = nalu
|
||
case codec.NAL_UNIT_SPS:
|
||
sps = nalu
|
||
vt.SPSInfo, _ = codec.ParseHevcSPS(nalu)
|
||
case codec.NAL_UNIT_PPS:
|
||
pps = nalu
|
||
}
|
||
}
|
||
if vps != nil && sps != nil && pps != nil {
|
||
extraData, err := codec.BuildH265SeqHeaderFromVpsSpsPps(vps, sps, pps)
|
||
if err != nil {
|
||
return
|
||
}
|
||
vt.ExtraData = &VideoPack{
|
||
NALUs: [][]byte{vps, sps, pps},
|
||
}
|
||
vt.ExtraData.Payload = extraData
|
||
}
|
||
if vt.ExtraData != nil {
|
||
vt.PushNalu = func(ts uint32, cts uint32, nalus ...[]byte) {
|
||
var nonIDRs [][]byte
|
||
for _, nalu := range nalus {
|
||
naluLen := len(nalu)
|
||
if naluLen == 0 {
|
||
continue
|
||
}
|
||
/*
|
||
0 1
|
||
0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5
|
||
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|
||
|F| Type | LayerId | TID |
|
||
+-------------+-----------------+
|
||
Forbidden zero(F) : 1 bit
|
||
NAL unit type(Type) : 6 bits
|
||
NUH layer ID(LayerId) : 6 bits
|
||
NUH temporal ID plus 1 (TID) : 3 bits
|
||
*/
|
||
|
||
naluType := nalu[0] & naluTypeBitmask_hevc >> 1
|
||
switch naluType {
|
||
case codec.NAL_UNIT_VPS:
|
||
vps = nalu
|
||
vt.ExtraData.NALUs[0] = vps
|
||
case codec.NAL_UNIT_SPS:
|
||
sps = nalu
|
||
vt.ExtraData.NALUs[1] = sps
|
||
vt.SPSInfo, _ = codec.ParseHevcSPS(nalu)
|
||
case codec.NAL_UNIT_PPS:
|
||
pps = nalu
|
||
vt.ExtraData.NALUs[2] = pps
|
||
extraData, err := codec.BuildH265SeqHeaderFromVpsSpsPps(vps, sps, pps)
|
||
if err != nil {
|
||
return
|
||
}
|
||
vt.ExtraData.Payload = extraData
|
||
case codec.NAL_UNIT_CODED_SLICE_BLA,
|
||
codec.NAL_UNIT_CODED_SLICE_BLANT,
|
||
codec.NAL_UNIT_CODED_SLICE_BLA_N_LP,
|
||
codec.NAL_UNIT_CODED_SLICE_IDR,
|
||
codec.NAL_UNIT_CODED_SLICE_IDR_N_LP,
|
||
codec.NAL_UNIT_CODED_SLICE_CRA:
|
||
vt.IDR = true
|
||
vt.SetTs(ts)
|
||
vt.CompositionTime = cts
|
||
vt.SetNalu0(nalu)
|
||
vt.addBytes(naluLen)
|
||
vt.push()
|
||
case 0, 1, 2, 3, 4, 5, 6, 7, 9:
|
||
nonIDRs = append(nonIDRs, nalu)
|
||
vt.addBytes(naluLen)
|
||
}
|
||
}
|
||
if len(nonIDRs) > 0 {
|
||
vt.IDR = false
|
||
vt.SetTs(ts)
|
||
vt.CompositionTime = cts
|
||
vt.NALUs = nonIDRs
|
||
vt.push()
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
vt.PushNalu(ts, cts, nalus...)
|
||
}
|
||
|
||
func (vt *VideoTrack) setCurrent() {
|
||
vt.AVTrack.setCurrent()
|
||
vt.VideoPack = vt.Value.(*VideoPack)
|
||
}
|
||
|
||
func (vt *VideoTrack) PushByteStream(ts uint32, payload []byte) {
|
||
if payload[1] == 0 {
|
||
vt.CodecID = payload[0] & 0x0F
|
||
switch vt.CodecID {
|
||
case 7:
|
||
var info codec.AVCDecoderConfigurationRecord
|
||
if _, err := info.Unmarshal(payload[5:]); err == nil {
|
||
vt.SPSInfo, _ = codec.ParseSPS(info.SequenceParameterSetNALUnit)
|
||
vt.nalulenSize = int(info.LengthSizeMinusOne&3 + 1)
|
||
vt.ExtraData = &VideoPack{
|
||
NALUs: [][]byte{info.SequenceParameterSetNALUnit, info.PictureParameterSetNALUnit},
|
||
}
|
||
vt.ExtraData.Payload = payload
|
||
vt.Stream.VideoTracks.AddTrack("h264", vt)
|
||
}
|
||
case 12:
|
||
if vps, sps, pps, err := codec.ParseVpsSpsPpsFromSeqHeaderWithoutMalloc(payload); err == nil {
|
||
vt.SPSInfo, _ = codec.ParseSPS(sps)
|
||
vt.nalulenSize = int(payload[26]) & 0x03
|
||
vt.ExtraData = &VideoPack{
|
||
NALUs: [][]byte{vps, sps, pps},
|
||
}
|
||
vt.ExtraData.Payload = payload
|
||
vt.Stream.VideoTracks.AddTrack("h265", vt)
|
||
}
|
||
}
|
||
} else {
|
||
if len(payload) < 4 {
|
||
return
|
||
}
|
||
vt.addBytes(len(payload))
|
||
vt.IDR = payload[0]>>4 == 1
|
||
vt.SetTs(ts)
|
||
vt.Payload = payload
|
||
vt.CompositionTime = utils.BigEndian.Uint24(payload[2:])
|
||
vt.ResetNALUs()
|
||
for nalus := payload[5:]; len(nalus) > vt.nalulenSize; {
|
||
nalulen := 0
|
||
for i := 0; i < vt.nalulenSize; i++ {
|
||
nalulen += int(nalus[i]) << (8 * (vt.nalulenSize - i - 1))
|
||
}
|
||
vt.NALUs = append(vt.NALUs, nalus[vt.nalulenSize:nalulen+vt.nalulenSize])
|
||
nalus = nalus[nalulen+vt.nalulenSize:]
|
||
}
|
||
vt.push()
|
||
}
|
||
}
|
||
|
||
func (vt *VideoTrack) push() {
|
||
if vt.Stream != nil {
|
||
vt.Stream.Update()
|
||
}
|
||
if vt.writeByteStream != nil {
|
||
vt.writeByteStream()
|
||
}
|
||
if vt.GetBPS(); vt.IDR {
|
||
vt.revIDR()
|
||
}
|
||
if nextPack := vt.NextValue().(*VideoPack); nextPack.IDR {
|
||
if vt.idrCount == 1 {
|
||
exRing := ring.New(5)
|
||
for x := exRing; x.Value == nil; x = x.Next() {
|
||
x.Value = &AVItem{Value: new(VideoPack)}
|
||
}
|
||
vt.Link(exRing) // 扩大缓冲环
|
||
} else {
|
||
vt.idrCount--
|
||
}
|
||
}
|
||
vt.Step()
|
||
vt.setCurrent()
|
||
}
|
||
|
||
func (vt *VideoTrack) Play(onVideo func(uint32, *VideoPack), exit1, exit2 <-chan struct{}) {
|
||
select {
|
||
case <-vt.WaitIDR:
|
||
case <-exit1:
|
||
return
|
||
case <-exit2: //可能等不到关键帧就退出了
|
||
return
|
||
}
|
||
vr := vt.SubRing(vt.IDRing) //从关键帧开始读取，首屏秒开
|
||
realSt := vt.PreItem().Timestamp // 当前时间戳
|
||
item, vp := vr.Read()
|
||
startTimestamp := item.Timestamp
|
||
for chase := true; ; item, vp = vr.Read() {
|
||
select {
|
||
case <-exit1:
|
||
return
|
||
case <-exit2:
|
||
return
|
||
default:
|
||
onVideo(item.Timestamp-startTimestamp, vp.(*VideoPack))
|
||
if chase {
|
||
if startTimestamp < realSt-10 {
|
||
startTimestamp += 10
|
||
} else {
|
||
startTimestamp = realSt
|
||
chase = false
|
||
}
|
||
}
|
||
vr.MoveNext()
|
||
}
|
||
}
|
||
}
|