Merge pull request #102 from joestarzxh/master

[fix]ps解析没有psm时，猜测视频类型，参考media-server

gb28181/mpegps/ps_demuxer.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"github.com/q191201771/lal/pkg/avc"
 	"github.com/q191201771/lal/pkg/hevc"
+	"github.com/q191201771/naza/pkg/nazalog"
 )
 
 type psStream struct {
@@ -26,6 +27,10 @@ func (p *psStream) setCid(cid PsStreamType) {
 	p.cid = cid
 }
 
+func (p *psStream) clearBuf() {
+	p.streamBuf = p.streamBuf[:0]
+}
+
 type PsDemuxer struct {
 	streamMap map[uint8]*psStream
 	pkg       *PsPacket
@@ -36,6 +41,10 @@ type PsDemuxer struct {
 	//decodeResult 解码ps包时的产生的错误
 	//这个回调主要用于debug，查看是否ps包存在问题
 	OnPacket func(pkg Display, decodeResult error)
+
+	verifyBuf []byte
+
+	log nazalog.Logger
 }
 
 func NewPsDemuxer() *PsDemuxer {
@@ -46,12 +55,6 @@ func NewPsDemuxer() *PsDemuxer {
 		OnFrame:   nil,
 		OnPacket:  nil,
 	}
-	//兼容没有发送psm的ps包
-	//兼容没有发送psm的ps包
-	streamH264 := newPsStream(uint8(PesStreamVideo), PsStreamH264)
-	streamG711A := newPsStream(uint8(PesStreamAudio), PsStreamG711A)
-	psDemuxer.streamMap[streamH264.sid] = streamH264
-	psDemuxer.streamMap[streamG711A.sid] = streamG711A
 	return psDemuxer
 }
 
@@ -163,6 +166,31 @@ func (psDemuxer *PsDemuxer) Input(data []byte) error {
 							stream.streamBuf = append(stream.streamBuf, psDemuxer.pkg.Pes.PesPayload...)
 							stream.pts = psDemuxer.pkg.Pes.Pts
 							stream.dts = psDemuxer.pkg.Pes.Dts
+						} else if psDemuxer.pkg.Pes.StreamId == uint8(PesStreamVideo) {
+							if len(psDemuxer.verifyBuf) > 256 {
+								psDemuxer.verifyBuf = psDemuxer.verifyBuf[:0]
+							}
+							psDemuxer.verifyBuf = append(psDemuxer.verifyBuf, psDemuxer.pkg.Pes.PesPayload...)
+							if h26x, err := mpegH26xVerify(psDemuxer.verifyBuf); err == nil {
+								switch h26x {
+								case CodecUnknown:
+								case CodecH264:
+									streamH264 := newPsStream(uint8(PesStreamVideo), PsStreamH264)
+									psDemuxer.streamMap[streamH264.sid] = streamH264
+									psDemuxer.demuxPespacket(streamH264, psDemuxer.pkg.Pes)
+								case CodecH265:
+									streamH265 := newPsStream(uint8(PesStreamVideo), PsStreamH265)
+									psDemuxer.streamMap[streamH265.sid] = streamH265
+									psDemuxer.demuxPespacket(streamH265, psDemuxer.pkg.Pes)
+								}
+							}
+						} else if psDemuxer.pkg.Pes.StreamId == uint8(PesStreamAudio) {
+							if _, found = psDemuxer.streamMap[uint8(PesStreamVideo)]; found {
+								psStreamType := audioVerify(psDemuxer.pkg.Pes.PesPayload)
+								streamAudio := newPsStream(uint8(PesStreamAudio), psStreamType)
+								psDemuxer.streamMap[streamAudio.sid] = streamAudio
+								psDemuxer.demuxPespacket(streamAudio, psDemuxer.pkg.Pes)
+							}
 						}
 					}
 				}
@@ -192,40 +220,18 @@ func (psDemuxer *PsDemuxer) Flush() {
 
 func (psDemuxer *PsDemuxer) guessCodecid(stream *psStream) {
 	if stream.sid&0xE0 == uint8(PesStreamAudio) {
-		stream.cid = PsStreamAac
+		psStreamType := audioVerify(psDemuxer.pkg.Pes.PesPayload)
+		stream.cid = psStreamType
 	} else if stream.sid&0xE0 == uint8(PesStreamVideo) {
-		h264score := 0
-		h265score := 0
-		SplitFrame(stream.streamBuf, func(nalu []byte) bool {
-			h264nalutype := avc.ParseNaluType(nalu[0])
-			h265nalutype := hevc.ParseNaluType(nalu[0])
-			if h264nalutype == avc.NaluTypeSps ||
-				h264nalutype == avc.NaluTypePps ||
-				h264nalutype == avc.NaluTypeIdrSlice {
-				h264score += 2
-			} else if h264nalutype < 5 {
-				h264score += 1
-			} else if h264nalutype > 20 {
-				h264score -= 1
-			}
-
-			if h265nalutype == hevc.NaluTypeSps ||
-				h265nalutype == hevc.NaluTypePps ||
-				h265nalutype == hevc.NaluTypeVps ||
-				(h265nalutype >= hevc.NaluTypeSliceBlaWlp && h265nalutype <= hevc.NaluTypeSliceRsvIrapVcl23) {
-				h265score += 2
-			} else if h265nalutype >= hevc.NaluTypeSliceTrailN && h265nalutype <= hevc.NaluTypeSliceRaslR {
-				h265score += 1
-			} else if h265nalutype > 40 {
-				h265score -= 1
-			}
-			if h264score > h265score && h264score >= 4 {
+		if h26x, err := mpegH26xVerify(stream.streamBuf); err == nil {
+			switch h26x {
+			case CodecUnknown:
+			case CodecH264:
 				stream.cid = PsStreamH264
-			} else if h264score < h265score && h265score >= 4 {
+			case CodecH265:
 				stream.cid = PsStreamH265
 			}
-			return true
-		})
+		}
 	}
 }
 

gb28181/mpegps/util.go

@@ -1,6 +1,17 @@
 package mpegps
 
-import "bytes"
+import (
+	"bytes"
+	"errors"
+)
+
+const (
+	CodecUnknown = iota
+	CodecH264
+	CodecH265
+	CodecH266
+	CodecMpeg4
+)
 
 var crc32table [256]uint32 = [256]uint32{
 	0x00000000, 0xB71DC104, 0x6E3B8209, 0xD926430D, 0xDC760413, 0x6B6BC517,
@@ -102,3 +113,85 @@ func H265NaluType(h265 []byte) uint8 {
 	loc, sc := FindStartCode(h265, 0)
 	return (h265[loc+int(sc)] >> 1) & 0x3F
 }
+
+func mpegH264FindNALU(data []byte) (int, int, error) {
+	var zeros, i int
+
+	for i = 0; i+2 < len(data); i++ {
+		if data[i] == 0x01 && zeros >= 2 {
+			return i + 1, zeros + 1, nil // 返回 NALU 的长度和前导零的数量
+		}
+
+		if data[i] == 0 {
+			zeros++
+		} else {
+			zeros = 0
+		}
+	}
+
+	return -1, 0, errors.New("no valid NALU found")
+}
+
+// 来自media-server
+func mpegH26xVerify(data []byte) (int, error) {
+	h264Flags := uint32(0x01A0)      // SPS/PPS/IDR
+	h265Flags := uint64(0x700000000) // VPS/SPS/PPS
+	h266Flags := uint32(0xC000)      // VPS/SPS/PPS
+	count := 0
+	h26x := [5][10]int{}
+
+	p := 0
+	end := len(data)
+
+	for p < end && count < len(h26x[0]) {
+		n, _, err := mpegH264FindNALU(data[p:])
+		if err != nil {
+			break
+		}
+		if p+n+1 > end {
+			break
+		}
+
+		h26x[0][count] = int(data[p+n]) & 0x1F          // H.264 NALU type
+		h26x[1][count] = (int(data[p+n]) >> 1) & 0x3F   // H.265 NALU type
+		h26x[2][count] = (int(data[p+n+1]) >> 3) & 0x1F // H.266 NALU type
+		h26x[3][count] = int(data[p+n])                 // MPEG-4 VOP start code
+		h26x[4][count] = int(data[p+n+1])               // MPEG-4 VOP coding type
+		count++
+
+		p += n // 移动到下一个 NALU
+	}
+
+	for n := 0; n < count; n++ {
+		h264Flags &= ^(1 << h26x[0][n])
+		h265Flags &= ^(1 << h26x[1][n])
+		h266Flags &= ^(1 << h26x[2][n])
+	}
+
+	if h264Flags == 0 && h265Flags != 0 && h266Flags != 0 {
+		// match SPS/PPS/IDR
+		return CodecH264, nil
+	} else if h265Flags == 0 && h264Flags != 0 && h266Flags != 0 {
+		// match VPS/SPS/PPS
+		return CodecH265, nil
+	} else if h266Flags == 0 && h264Flags != 0 && h265Flags != 0 {
+		// match SPS/PPS
+		return CodecH266, nil
+	} else if h26x[3][0] == 0xB0 && (h26x[4][0]&0x30) == 0 {
+		// match VOP start code
+		return CodecMpeg4, nil
+	}
+
+	return CodecUnknown, nil
+}
+func audioVerify(data []byte) PsStreamType {
+	if data[0] == 0xFF && (data[1]&0xF0) == 0xF0 && len(data) > 7 {
+		aacLen := ((int(data[3]) & 0x03) << 11) |
+			(int(data[4]) << 3) |
+			(int(data[5]) >> 5 & 0x07)
+		if len(data) == aacLen {
+			return PsStreamAac
+		}
+	}
+	return PsStreamG711A
+}