mkv in progress

format/mkv/demuxer.go

@@ -0,0 +1,147 @@
+package mkv
+
+import (
+	"encoding/binary"
+	"errors"
+	"io"
+	"time"
+
+	"github.com/deepch/vdk/av"
+	"github.com/deepch/vdk/codec/h264parser"
+	"github.com/deepch/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
+}