mirror of
https://github.com/aler9/gortsplib
synced 2025-10-05 07:06:58 +08:00
311 lines
7.4 KiB
Go
311 lines
7.4 KiB
Go
package h264
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"math"
|
|
"time"
|
|
|
|
"github.com/aler9/gortsplib/pkg/bits"
|
|
)
|
|
|
|
func getPOC(buf []byte, sps *SPS) (uint32, error) {
|
|
buf = AntiCompetitionRemove(buf[:6])
|
|
|
|
isIDR := NALUType(buf[0]&0x1F) == NALUTypeIDR
|
|
|
|
buf = buf[1:]
|
|
pos := 0
|
|
|
|
// first_mb_in_slice
|
|
_, err := bits.ReadGolombUnsigned(buf, &pos)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
// slice_type
|
|
_, err = bits.ReadGolombUnsigned(buf, &pos)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
// pic_parameter_set_id
|
|
_, err = bits.ReadGolombUnsigned(buf, &pos)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
// frame_num
|
|
_, err = bits.ReadBits(buf, &pos, int(sps.Log2MaxFrameNumMinus4+4))
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
if !sps.FrameMbsOnlyFlag {
|
|
return 0, fmt.Errorf("unsupported")
|
|
}
|
|
|
|
if isIDR {
|
|
// idr_pic_id
|
|
_, err = bits.ReadGolombUnsigned(buf, &pos)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
|
|
var picOrderCntLsb uint64
|
|
switch {
|
|
case sps.PicOrderCntType == 0:
|
|
picOrderCntLsb, err = bits.ReadBits(buf, &pos, int(sps.Log2MaxPicOrderCntLsbMinus4+4))
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
default:
|
|
return 0, fmt.Errorf("pic_order_cnt_type = 1 is unsupported")
|
|
}
|
|
|
|
return uint32(picOrderCntLsb), nil
|
|
}
|
|
|
|
func findPOC(nalus [][]byte, sps *SPS) (uint32, error) {
|
|
for _, nalu := range nalus {
|
|
typ := NALUType(nalu[0] & 0x1F)
|
|
if typ == NALUTypeIDR || typ == NALUTypeNonIDR {
|
|
poc, err := getPOC(nalu, sps)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return poc, nil
|
|
}
|
|
}
|
|
return 0, fmt.Errorf("POC not found")
|
|
}
|
|
|
|
func getPOCDiff(poc1 uint32, poc2 uint32, sps *SPS) int32 {
|
|
diff := int32(poc1) - int32(poc2)
|
|
switch {
|
|
case diff < -((1 << (sps.Log2MaxPicOrderCntLsbMinus4 + 3)) - 1):
|
|
diff += (1 << (sps.Log2MaxPicOrderCntLsbMinus4 + 4))
|
|
|
|
case diff > ((1 << (sps.Log2MaxPicOrderCntLsbMinus4 + 3)) - 1):
|
|
diff -= (1 << (sps.Log2MaxPicOrderCntLsbMinus4 + 4))
|
|
}
|
|
return diff
|
|
}
|
|
|
|
func getSEIPicTimingDPBOutputDelay(buf []byte, sps *SPS) (uint32, bool) {
|
|
buf = AntiCompetitionRemove(buf)
|
|
pos := 1
|
|
|
|
for {
|
|
if pos >= (len(buf) - 1) {
|
|
return 0, false
|
|
}
|
|
|
|
payloadType := 0
|
|
for {
|
|
byt := buf[pos]
|
|
pos++
|
|
payloadType += int(byt)
|
|
if byt != 0xFF {
|
|
break
|
|
}
|
|
}
|
|
|
|
payloadSize := 0
|
|
for {
|
|
byt := buf[pos]
|
|
pos++
|
|
payloadSize += int(byt)
|
|
if byt != 0xFF {
|
|
break
|
|
}
|
|
}
|
|
|
|
if payloadType == 1 { // timing info
|
|
buf2 := buf[pos:]
|
|
pos2 := 0
|
|
|
|
// cpbRemovalDelay
|
|
_, err := bits.ReadBits(buf2, &pos2, int(sps.VUI.NalHRD.CpbRemovalDelayLengthMinus1+1))
|
|
if err != nil {
|
|
return 0, false
|
|
}
|
|
|
|
tmp, err := bits.ReadBits(buf2, &pos2, int(sps.VUI.NalHRD.DpbOutputDelayLengthMinus1+1))
|
|
if err != nil {
|
|
return 0, false
|
|
}
|
|
dpbOutputDelay := uint32(tmp)
|
|
|
|
return dpbOutputDelay, true
|
|
}
|
|
|
|
pos += payloadSize
|
|
}
|
|
}
|
|
|
|
func findSEIPicTimingDPBOutputDelay(nalus [][]byte, sps *SPS) (uint32, bool) {
|
|
for _, nalu := range nalus {
|
|
typ := NALUType(nalu[0] & 0x1F)
|
|
if typ == NALUTypeSEI {
|
|
ret, ok := getSEIPicTimingDPBOutputDelay(nalu, sps)
|
|
if ok {
|
|
return ret, true
|
|
}
|
|
}
|
|
}
|
|
return 0, false
|
|
}
|
|
|
|
// DTSExtractor is a utility that allows to extract NALU DTS from PTS.
|
|
type DTSExtractor struct {
|
|
sps []byte
|
|
spsp *SPS
|
|
prevPTS time.Duration
|
|
prevDTS *time.Duration
|
|
prevPOCDiff int32
|
|
expectedPOC uint32
|
|
ptsDTSOffset time.Duration
|
|
firstDPBOutputDelay *uint32
|
|
}
|
|
|
|
// NewDTSExtractor allocates a DTSExtractor.
|
|
func NewDTSExtractor() *DTSExtractor {
|
|
return &DTSExtractor{}
|
|
}
|
|
|
|
func (d *DTSExtractor) extractInner(nalus [][]byte, pts time.Duration) (time.Duration, int32, error) {
|
|
idrPresent := false
|
|
|
|
for _, nalu := range nalus {
|
|
typ := NALUType(nalu[0] & 0x1F)
|
|
switch typ {
|
|
case NALUTypeSPS:
|
|
if d.sps == nil || !bytes.Equal(d.sps, nalu) {
|
|
var spsp SPS
|
|
err := spsp.Unmarshal(nalu)
|
|
if err != nil {
|
|
return 0, 0, fmt.Errorf("invalid SPS: %v", err)
|
|
}
|
|
d.sps = append([]byte(nil), nalu...)
|
|
d.spsp = &spsp
|
|
|
|
if d.spsp.VUI != nil && d.spsp.VUI.TimingInfo != nil &&
|
|
d.spsp.VUI.BitstreamRestriction != nil {
|
|
d.ptsDTSOffset = time.Duration(math.Round(float64(
|
|
time.Duration(d.spsp.VUI.BitstreamRestriction.MaxNumReorderFrames)*time.Second*
|
|
time.Duration(d.spsp.VUI.TimingInfo.NumUnitsInTick)*2) / float64(d.spsp.VUI.TimingInfo.TimeScale)))
|
|
} else {
|
|
d.ptsDTSOffset = 0
|
|
}
|
|
}
|
|
|
|
case NALUTypeIDR:
|
|
idrPresent = true
|
|
}
|
|
}
|
|
|
|
if d.spsp == nil {
|
|
return 0, 0, fmt.Errorf("SPS not received yet")
|
|
}
|
|
|
|
switch {
|
|
// DTS is computed by using POC, timing infos and max_num_reorder_frames
|
|
case d.spsp.PicOrderCntType != 2 &&
|
|
d.spsp.VUI != nil && d.spsp.VUI.TimingInfo != nil && d.spsp.VUI.BitstreamRestriction != nil:
|
|
if idrPresent {
|
|
d.expectedPOC = 0
|
|
return pts - d.ptsDTSOffset, 0, nil
|
|
}
|
|
|
|
// compute expectedPOC immediately in order to store it even in case of errors
|
|
d.expectedPOC += 2
|
|
d.expectedPOC &= ((1 << (d.spsp.Log2MaxPicOrderCntLsbMinus4 + 4)) - 1)
|
|
|
|
poc, err := findPOC(nalus, d.spsp)
|
|
if err != nil {
|
|
return 0, 0, err
|
|
}
|
|
|
|
pocDiff := getPOCDiff(poc, d.expectedPOC, d.spsp)
|
|
|
|
if pocDiff == 0 {
|
|
return pts - d.ptsDTSOffset, 0, nil
|
|
}
|
|
|
|
// special case to eliminate errors near 0
|
|
if pocDiff == -int32(d.spsp.VUI.BitstreamRestriction.MaxNumReorderFrames)*2 {
|
|
return pts, pocDiff, nil
|
|
}
|
|
|
|
if d.prevPOCDiff == 0 {
|
|
if pocDiff == -2 {
|
|
return 0, 0, fmt.Errorf("invalid frame POC")
|
|
}
|
|
|
|
return d.prevPTS - d.ptsDTSOffset +
|
|
time.Duration(math.Round(float64(pts-d.prevPTS)/float64(pocDiff/2+1))), pocDiff, nil
|
|
}
|
|
|
|
// pocDiff : prevPOCDiff = (pts - dts - ptsDTSOffset) : (prevPTS - prevDTS - ptsDTSOffset)
|
|
return pts - d.ptsDTSOffset + time.Duration(math.Round(float64(*d.prevDTS-d.prevPTS+d.ptsDTSOffset)*
|
|
float64(pocDiff)/float64(d.prevPOCDiff))), pocDiff, nil
|
|
|
|
// DTS is computed from SEI
|
|
case d.spsp.VUI != nil && d.spsp.VUI.TimingInfo != nil && d.spsp.VUI.NalHRD != nil:
|
|
dpbOutputDelay, ok := findSEIPicTimingDPBOutputDelay(nalus, d.spsp)
|
|
if !ok {
|
|
// some streams declare that they use SEI pic timings, but they don't.
|
|
// assume PTS = DTS.
|
|
return pts, 0, nil
|
|
}
|
|
|
|
// workaround for nvenc
|
|
// nvenc puts a wrong dpbOutputDelay in IDR frames that follows the first one.
|
|
// save the dpbOutputDelay of the first IDR frame and use if for subsequent
|
|
// IDR frames.
|
|
// https://forums.developer.nvidia.com/t/nvcodec-h-264-encoder-sei-pic-timing-dpb-output-delay/156050
|
|
// https://forums.developer.nvidia.com/t/h264-pic-timing-sei-message/71188
|
|
if idrPresent &&
|
|
d.spsp.VUI.NalHRD.CpbRemovalDelayLengthMinus1 == 15 &&
|
|
d.spsp.VUI.NalHRD.DpbOutputDelayLengthMinus1 == 5 {
|
|
if d.firstDPBOutputDelay == nil {
|
|
d.firstDPBOutputDelay = &dpbOutputDelay
|
|
} else {
|
|
dpbOutputDelay = *d.firstDPBOutputDelay
|
|
}
|
|
}
|
|
|
|
return pts - time.Duration(dpbOutputDelay)/2*time.Second*
|
|
time.Duration(d.spsp.VUI.TimingInfo.NumUnitsInTick)*2/time.Duration(d.spsp.VUI.TimingInfo.TimeScale), 0, nil
|
|
|
|
// we assume PTS = DTS
|
|
default:
|
|
return pts, 0, nil
|
|
}
|
|
}
|
|
|
|
// Extract extracts the DTS of a group of NALUs.
|
|
func (d *DTSExtractor) Extract(nalus [][]byte, pts time.Duration) (time.Duration, error) {
|
|
dts, pocDiff, err := d.extractInner(nalus, pts)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
if dts > pts {
|
|
return 0, fmt.Errorf("DTS is greater than PTS")
|
|
}
|
|
|
|
if d.prevDTS != nil && dts <= *d.prevDTS {
|
|
return 0, fmt.Errorf("DTS is not monotonically increasing, was %v, now is %v",
|
|
*d.prevDTS, dts)
|
|
}
|
|
|
|
d.prevPTS = pts
|
|
d.prevDTS = &dts
|
|
d.prevPOCDiff = pocDiff
|
|
return dts, err
|
|
}
|