148 lines
2.9 KiB
Go
148 lines
2.9 KiB
Go
package mkv
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"errors"
|
|
"io"
|
|
"time"
|
|
|
|
"git.r-2.top/kunmeng/vdk/av"
|
|
"git.r-2.top/kunmeng/vdk/codec/h264parser"
|
|
"git.r-2.top/kunmeng/vdk/format/mkv/mkvio"
|
|
)
|
|
|
|
type Demuxer struct {
|
|
r *mkvio.Document
|
|
pkts []av.Packet
|
|
sps []byte
|
|
pps []byte
|
|
streams []*Stream
|
|
ps uint32
|
|
stage int
|
|
fc int
|
|
ls time.Duration
|
|
}
|
|
|
|
func NewDemuxer(r io.Reader) *Demuxer {
|
|
|
|
return &Demuxer{
|
|
r: mkvio.InitDocument(r),
|
|
}
|
|
|
|
}
|
|
|
|
func (self *Demuxer) Streams() (streams []av.CodecData, err error) {
|
|
|
|
if err = self.probe(); err != nil {
|
|
return
|
|
}
|
|
|
|
for _, stream := range self.streams {
|
|
streams = append(streams, stream.CodecData)
|
|
}
|
|
|
|
if len(streams) == 0 {
|
|
return nil, errors.New("streams not found")
|
|
}
|
|
|
|
return streams, err
|
|
|
|
}
|
|
|
|
func (self *Demuxer) probe() (err error) {
|
|
if self.stage == 0 {
|
|
|
|
var el *mkvio.Element
|
|
el, err = self.r.GetVideoCodec()
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
if el.ElementRegister.ID == mkvio.ElementCodecPrivate.ID {
|
|
payload := el.Content[6:]
|
|
var reader int
|
|
for pos := 0; pos < len(payload); pos = reader {
|
|
lens := int(binary.BigEndian.Uint16(payload[reader:]))
|
|
reader += 2
|
|
nal := payload[reader : reader+lens]
|
|
naluType := nal[0] & 0x1f
|
|
switch naluType {
|
|
case h264parser.NALU_SPS:
|
|
self.sps = nal
|
|
case h264parser.NALU_PPS:
|
|
self.pps = nal
|
|
}
|
|
reader += lens
|
|
reader++
|
|
}
|
|
}
|
|
|
|
if len(self.sps) > 0 && len(self.pps) > 0 {
|
|
var codec av.CodecData
|
|
codec, err = h264parser.NewCodecDataFromSPSAndPPS(self.sps, self.pps)
|
|
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
stream := &Stream{}
|
|
stream.idx = 0
|
|
stream.demuxer = self
|
|
stream.CodecData = codec
|
|
self.streams = append(self.streams, stream)
|
|
|
|
}
|
|
self.stage++
|
|
}
|
|
return
|
|
}
|
|
|
|
func (self *Demuxer) ReadPacket() (pkt av.Packet, err error) {
|
|
|
|
var el mkvio.Element
|
|
|
|
for {
|
|
el, err = self.r.ParseElement()
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
if el.Type == 6 && el.ElementRegister.ID == mkvio.ElementSimpleBlock.ID {
|
|
self.fc++
|
|
nals, _ := h264parser.SplitNALUs(el.Content[4:])
|
|
for _, nal := range nals {
|
|
|
|
naluType := nal[0] & 0x1f
|
|
|
|
if naluType == 5 {
|
|
l1 := int(binary.BigEndian.Uint16(el.Content[2:4]))
|
|
dur := time.Duration(uint32(l1)) * time.Millisecond
|
|
self.ls += time.Duration(uint32(l1)) * time.Millisecond
|
|
self.ps = 0
|
|
pkt = av.Packet{IsKeyFrame: true, Idx: 0, Duration: dur, Time: self.ls, Data: append(binSize(len(nal)), nal...)}
|
|
return
|
|
|
|
} else if naluType == 1 {
|
|
|
|
l1 := int(binary.BigEndian.Uint16(el.Content[1:3]))
|
|
dur := time.Duration(uint32(l1)-self.ps) * time.Millisecond
|
|
self.ls += time.Duration(uint32(l1)-self.ps) * time.Millisecond
|
|
self.ps = uint32(l1)
|
|
pkt = av.Packet{Idx: 0, Duration: dur, Time: self.ls, Data: append(binSize(len(nal)), nal...)}
|
|
return
|
|
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return
|
|
|
|
}
|
|
|
|
func binSize(val int) []byte {
|
|
buf := make([]byte, 4)
|
|
binary.BigEndian.PutUint32(buf, uint32(val))
|
|
return buf
|
|
}
|